This is an automated email from the ASF dual-hosted git repository. krickert pushed a commit to branch sentencepiece in repository https://gitbox.apache.org/repos/asf/opennlp.git
commit 7f20e7663d7fb61d8ad4632efa09feef6ad47a71 Author: Kristian Rickert <[email protected]> AuthorDate: Thu Jul 16 06:27:30 2026 -0400 OPENNLP-1885: Make the tokenizer graph serializable with computed UIDs, name the format constants, document every helper --- .../opennlp/tools/tokenize/WordpieceEncoder.java | 77 ++++++++++-- .../test/java/opennlp/dl/CreateTokenizerTest.java | 3 +- .../tools/tokenize/ReferenceBertPipeline.java | 2 + .../WordpieceEncoderReferenceSequencesTest.java | 130 ++++++++------------- .../tools/tokenize/WordpieceEncoderTest.java | 47 +++++--- opennlp-docs/src/docbkx/tokenizer.xml | 4 +- .../opennlp/subword/sentencepiece/BpeEncoder.java | 40 ++++--- .../opennlp/subword/sentencepiece/ByteBuilder.java | 8 +- .../subword/sentencepiece/DoubleArrayTrie.java | 32 ++++- .../opennlp/subword/sentencepiece/IntBuilder.java | 8 +- .../subword/sentencepiece/ModelProtoReader.java | 115 +++++++++++++++--- .../opennlp/subword/sentencepiece/PieceTrie.java | 68 +++++++++-- .../sentencepiece/SentencePieceNormalizer.java | 41 +++---- .../sentencepiece/SentencePieceTokenizer.java | 44 +++++-- .../subword/sentencepiece/UnigramEncoder.java | 15 ++- .../opennlp/subword/sentencepiece/Utf8Text.java | 8 ++ .../sentencepiece/SentencePieceFixtures.java | 125 ++++++++++++++++++++ .../SentencePieceModelValidationTest.java | 6 +- .../sentencepiece/SentencePieceParityTest.java | 65 +---------- .../SentencePieceRealModelEvalTest.java | 50 +------- .../SentencePieceTokenizerSerializationTest.java | 71 +++++++++++ 21 files changed, 653 insertions(+), 306 deletions(-) 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 ab016bbaa..b130f45dc 100644 --- a/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceEncoder.java +++ b/opennlp-api/src/main/java/opennlp/tools/tokenize/WordpieceEncoder.java @@ -40,7 +40,8 @@ import java.util.Set; * <p>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.</p> + * 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.</p> * * <p>Instances are immutable and safe for concurrent use by multiple threads.</p> * @@ -48,6 +49,10 @@ import java.util.Set; */ public final class WordpieceEncoder implements SubwordTokenizer { + // The wordpiece vocabulary convention: a piece with this prefix continues the current word, + // so it can only match after the word's first piece. + private static final String CONTINUATION_PREFIX = "##"; + // The reference implementation's limit: longer words become the unknown piece. private static final int MAX_WORD_CHARACTERS = 100; @@ -144,6 +149,15 @@ public final class WordpieceEncoder implements SubwordTokenizer { this.unknownId = requiredId(byPiece, unknownToken); } + /** + * Converts an ordered vocabulary list into the piece-to-id mapping, assigning each piece its + * index as the id. + * + * @param vocabulary The ordered vocabulary. + * @return The piece-to-id mapping. + * @throws IllegalArgumentException Thrown if the list is null or contains a null or duplicate + * entry. + */ private static Map<String, Integer> byPiece(List<String> vocabulary) { if (vocabulary == null) { throw new IllegalArgumentException("The vocabulary must not be null."); @@ -162,6 +176,14 @@ public final class WordpieceEncoder implements SubwordTokenizer { 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<String, Integer> ids, String specialToken) { final Integer id = ids.get(specialToken); if (id == null) { @@ -171,11 +193,7 @@ public final class WordpieceEncoder implements SubwordTokenizer { return id; } - /** - * {@inheritDoc} - * - * @throws IllegalArgumentException Thrown if {@code text} is null. - */ + /** {@inheritDoc} */ @Override public List<SubwordPiece> encode(CharSequence text) { if (text == null) { @@ -236,7 +254,7 @@ public final class WordpieceEncoder implements SubwordTokenizer { while (start < end) { String substring = new String(mapped.chars, start, end - start); if (start > from) { - substring = "##" + substring; + substring = CONTINUATION_PREFIX + substring; } if (vocabulary.contains(substring)) { wordPieces.add(new SubwordPiece(substring, ids.get(substring), @@ -268,12 +286,24 @@ public final class WordpieceEncoder implements SubwordTokenizer { private int[] ends; private int length; + /** + * Instantiates an empty mapped text. + * + * @param capacity The initial capacity hint in chars. + */ private MappedText(int capacity) { chars = new char[capacity]; starts = new int[capacity]; ends = new int[capacity]; } + /** + * Appends one char with the original-text range it came from. + * + * @param c The char to append. + * @param originalStart The inclusive original-text start of the char. + * @param originalEnd The exclusive original-text end of the char. + */ private void add(char c, int originalStart, int originalEnd) { if (length == chars.length) { final int capacity = Math.max(16, length * 2); @@ -287,6 +317,13 @@ public final class WordpieceEncoder implements SubwordTokenizer { length++; } + /** + * Appends every char of a string, all sharing one original-text range. + * + * @param s The string to append. + * @param originalStart The inclusive original-text start shared by all chars. + * @param originalEnd The exclusive original-text end shared by all chars. + */ private void add(String s, int originalStart, int originalEnd) { for (int i = 0; i < s.length(); i++) { add(s.charAt(i), originalStart, originalEnd); @@ -385,8 +422,20 @@ public final class WordpieceEncoder implements SubwordTokenizer { return out; } + /** + * Lower cases and accent-strips one non-space run, emitting per-character ranges when the + * transformation is reproducible per code point and the run's full range otherwise. + * + * @param in The input text with per-character ranges. + * @param from The inclusive start of the run in {@code in}. + * @param to The exclusive end of the run in {@code in}. + * @param out The output text to append to. + */ private static void transformRun(MappedText in, int from, int to, MappedText out) { final String run = new String(in.chars, from, to - from); + // Locale.ROOT lower casing is the reference behavior of BERT's do_lower_case: the reference + // pipeline applies the full locale-independent Unicode case mappings (including one-to-many + // ones like the dotted capital I), which a per-code-point mapping cannot reproduce. final String content = stripAccents(run.toLowerCase(Locale.ROOT)); // Rerun per code point to learn how many output chars each input code point produces. @@ -419,6 +468,13 @@ public final class WordpieceEncoder implements SubwordTokenizer { } } + /** + * 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()); @@ -430,6 +486,13 @@ public final class WordpieceEncoder implements SubwordTokenizer { return stripped.toString(); } + /** + * Reads the code point at an index, joining a surrogate pair when one starts there. + * + * @param text The text to read from. + * @param index The char index to read at. + * @return The code point at {@code index}. + */ private static int codePointAt(MappedText text, int index) { final char c = text.chars[index]; if (Character.isHighSurrogate(c) && index + 1 < text.length 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 5131ae84e..5ae7e3ebe 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 @@ -101,7 +101,8 @@ public class CreateTokenizerTest { final Map<String, Integer> vocab = robertaVocab(); vocab.remove(WordpieceTokenizer.ROBERTA_UNK_TOKEN); - assertThrows(IllegalArgumentException.class, () -> AbstractDL.createPipelineTokenizer(vocab, false)); + assertThrows(IllegalArgumentException.class, + () -> AbstractDL.createPipelineTokenizer(vocab, false)); assertThrows(IllegalArgumentException.class, () -> AbstractDL.createWordpieceTokenizer(vocab)); } diff --git a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/ReferenceBertPipeline.java b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/ReferenceBertPipeline.java index 867c28dbf..a07d6faaf 100644 --- a/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/ReferenceBertPipeline.java +++ b/opennlp-core/opennlp-runtime/src/test/java/opennlp/tools/tokenize/ReferenceBertPipeline.java @@ -47,6 +47,8 @@ final class ReferenceBertPipeline { String normalized = cleanText(text); normalized = isolateCjkCharacters(normalized); if (lowerCase) { + // Locale.ROOT lower casing is the reference behavior of BERT's do_lower_case: the full + // locale-independent Unicode case mappings, including one-to-many ones. normalized = stripAccents(normalized.toLowerCase(Locale.ROOT)); } return BertNormalization.isolatePunctuation(normalized); 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 68e8f219f..064d644a5 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 @@ -17,13 +17,17 @@ package opennlp.tools.tokenize; import java.util.List; +import java.util.stream.Stream; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; /** - * The reference token sequences of the removed full-pipeline {@code Tokenizer}, re-asserted - * against {@link WordpieceEncoder}. + * Reference token-sequence expectations for {@link WordpieceEncoder}, covering lower casing, + * accent stripping, punctuation and CJK isolation, and text cleaning. * <p> * All expected token sequences in this test were generated with the HuggingFace * {@code tokenizers} reference implementation ({@code BertWordPieceTokenizer}) @@ -43,92 +47,56 @@ public class WordpieceEncoderReferenceSequencesTest { "\u6211", "\u7231", // CJK "natural", "language", "processing"); - @Test - void testLowerCasesCapitalizedWords() { - final WordpieceEncoder encoder = new WordpieceEncoder(VOCABULARY); - final String[] tokens = - encoder.encodeToPieces("The quick brown fox jumps over the lazy dog."); - - final String[] expected = {"[CLS]", "the", "quick", "brown", "fox", "jumps", "over", - "the", "lazy", "dog", ".", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); - } - - @Test - void testLowerCasesBeforeWordpieceSplitting() { - final WordpieceEncoder encoder = new WordpieceEncoder(VOCABULARY); - final String[] tokens = encoder.encodeToPieces("Embeddings"); - - final String[] expected = {"[CLS]", "em", "##bed", "##ding", "##s", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); - } - - @Test - void testStripsAccentsButKeepsNonCombiningCharacters() { - final WordpieceEncoder encoder = new WordpieceEncoder(VOCABULARY); - // 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. - final String[] tokens = encoder.encodeToPieces("W\u00fcrttemberg Stra\u00dfe"); - - final String[] expected = {"[CLS]", "wurttemberg", "[UNK]", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); + /** + * The reference input and expected-sequence pairs, one argument set per pipeline behavior. + * + * @return The (input, expected pieces) pairs. + */ + static Stream<Arguments> referenceSequences() { + return Stream.of( + // Lower cases capitalized words. + Arguments.of("The quick brown fox jumps over the lazy dog.", + new String[] {"[CLS]", "the", "quick", "brown", "fox", "jumps", "over", + "the", "lazy", "dog", ".", "[SEP]"}), + // Lower cases before wordpiece splitting. + Arguments.of("Embeddings", + new String[] {"[CLS]", "em", "##bed", "##ding", "##s", "[SEP]"}), + // The u-umlaut decomposes to u plus a combining diaeresis and the mark is stripped; + // the sharp s is not a combining mark and must survive, leaving an OOV token. + Arguments.of("W\u00fcrttemberg Stra\u00dfe", + new String[] {"[CLS]", "wurttemberg", "[UNK]", "[SEP]"}), + // Splits punctuation runs into single characters. + Arguments.of("Wait... what?!", + new String[] {"[CLS]", "wait", ".", ".", ".", "what", "?", "!", "[SEP]"}), + // Splits apostrophes as punctuation. + Arguments.of("don't", + new String[] {"[CLS]", "don", "'", "t", "[SEP]"}), + // Isolates CJK ideographs into single-character pieces. + Arguments.of("\u6211\u7231natural language processing", + new String[] {"[CLS]", "\u6211", "\u7231", "natural", "language", + "processing", "[SEP]"}), + // Tab and no-break space are whitespace; the NUL character is removed, + // joining "brown" and "fox" into one out-of-vocabulary token. + Arguments.of("the\tquick\u00a0brown\u0000fox", + new String[] {"[CLS]", "the", "quick", "[UNK]", "[SEP]"}), + // The reference implementation treats all C* categories as control + // characters: private use (U+E000, Co) and noncharacters (U+FDD0, Cn) + // are removed, joining the surrounding text into one OOV token. + Arguments.of("fox\ue000jumps and fox\ufdd0jumps", + new String[] {"[CLS]", "[UNK]", "[UNK]", "[UNK]", "[SEP]"})); } - @Test - void testSplitsPunctuationRunsIntoSingleCharacters() { + @ParameterizedTest + @MethodSource("referenceSequences") + void testEncodesTheReferenceSequence(String input, String[] expected) { final WordpieceEncoder encoder = new WordpieceEncoder(VOCABULARY); - final String[] tokens = encoder.encodeToPieces("Wait... what?!"); - - final String[] expected = {"[CLS]", "wait", ".", ".", ".", "what", "?", "!", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); - } - - @Test - void testSplitsApostrophesAsPunctuation() { - final WordpieceEncoder encoder = new WordpieceEncoder(VOCABULARY); - final String[] tokens = encoder.encodeToPieces("don't"); - - final String[] expected = {"[CLS]", "don", "'", "t", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); - } - - @Test - void testIsolatesCjkIdeographs() { - final WordpieceEncoder encoder = new WordpieceEncoder(VOCABULARY); - final String[] tokens = encoder.encodeToPieces("\u6211\u7231natural language processing"); - - final String[] expected = {"[CLS]", "\u6211", "\u7231", "natural", "language", - "processing", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); - } - - @Test - void testCleansControlCharactersAndNormalizesWhitespace() { - final WordpieceEncoder encoder = new WordpieceEncoder(VOCABULARY); - // Tab and no-break space are whitespace; the NUL character is removed, - // joining "brown" and "fox" into one out-of-vocabulary token. - final String[] tokens = encoder.encodeToPieces("the\tquick\u00a0brown\u0000fox"); - - final String[] expected = {"[CLS]", "the", "quick", "[UNK]", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); - } - - @Test - void testRemovesPrivateUseAndUnassignedCharacters() { - final WordpieceEncoder encoder = new WordpieceEncoder(VOCABULARY); - // The reference implementation treats all C* categories as control - // characters: private use (U+E000, Co) and noncharacters (U+FDD0, Cn) - // are removed, joining the surrounding text into one OOV token. - final String[] tokens = encoder.encodeToPieces("fox\ue000jumps and fox\ufdd0jumps"); - - final String[] expected = {"[CLS]", "[UNK]", "[UNK]", "[UNK]", "[SEP]"}; - Assertions.assertArrayEquals(expected, tokens); + Assertions.assertArrayEquals(expected, encoder.encodeToPieces(input), + "sequence broke on: " + input); } @Test void testRejectsNullSpecialTokens() { - // The encoder's contract throws IllegalArgumentException where the removed class threw - // NullPointerException. + // 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, 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 610bba5ed..cf545b183 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 @@ -20,8 +20,12 @@ import java.util.ArrayList; import java.util.HashSet; import java.util.List; import java.util.Random; +import java.util.stream.Stream; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -55,11 +59,13 @@ class WordpieceEncoderTest { assertEquals(expectedEnd, piece.end(), "end of " + piece); } - @Test - void testPieceSequenceMatchesTheReferencePipelineOnCuratedInputs() { - final ReferenceBertPipeline reference = new ReferenceBertPipeline(new HashSet<>(VOCAB), true); - final WordpieceEncoder encoder = uncased(); - final String[] inputs = { + /** + * The curated parity inputs, each exercising a normalization step of the pipeline. + * + * @return The inputs. + */ + static Stream<String> curatedInputs() { + return Stream.of( "", " ", "Hello, WORLD!", @@ -81,12 +87,16 @@ class WordpieceEncoderTest { "\uD83D\uDE00", "!!!", "a".repeat(101), - "he said: \u00ABhello\u00BB.", - }; - for (final String input : inputs) { - assertArrayEquals(reference.tokenize(input), encoder.encodeToPieces(input), - "parity broke on: " + input); - } + "he said: \u00ABhello\u00BB."); + } + + @ParameterizedTest + @MethodSource("curatedInputs") + void testPieceSequenceMatchesTheReferencePipelineOnCuratedInputs(String input) { + final ReferenceBertPipeline reference = new ReferenceBertPipeline(new HashSet<>(VOCAB), true); + final WordpieceEncoder encoder = uncased(); + assertArrayEquals(reference.tokenize(input), encoder.encodeToPieces(input), + "parity broke on: " + input); } @Test @@ -193,14 +203,13 @@ class WordpieceEncoderTest { assertPiece(pieces.get(2), "hello", 4, 6, 11); } - @Test - void testEmptyAndBlankTextEncodeToTheFramePiecesOnly() { - for (final String input : new String[] {"", " "}) { - final List<SubwordPiece> pieces = uncased().encode(input); - assertEquals(2, pieces.size()); - assertPiece(pieces.get(0), "[CLS]", 2, 0, 0); - assertPiece(pieces.get(1), "[SEP]", 3, input.length(), input.length()); - } + @ParameterizedTest + @ValueSource(strings = {"", " "}) + void testEmptyAndBlankTextEncodeToTheFramePiecesOnly(String input) { + final List<SubwordPiece> pieces = uncased().encode(input); + assertEquals(2, pieces.size(), "frame pieces broke on <" + input + ">"); + assertPiece(pieces.get(0), "[CLS]", 2, 0, 0); + assertPiece(pieces.get(1), "[SEP]", 3, input.length(), input.length()); } @Test diff --git a/opennlp-docs/src/docbkx/tokenizer.xml b/opennlp-docs/src/docbkx/tokenizer.xml index 2c87a280e..3449cf0d1 100644 --- a/opennlp-docs/src/docbkx/tokenizer.xml +++ b/opennlp-docs/src/docbkx/tokenizer.xml @@ -546,8 +546,8 @@ wordTokenizer.tokenize("The quick brown fox.", (start, end, type) -> { 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 <code>encodeToIds</code> and <code>encodeToPieces</code> methods return just the ids or - the piece strings when the spans are not needed. Implementations are safe for concurrent - use by multiple threads. + the piece strings when the spans are not needed. Implementations are expected to be safe + for concurrent use by multiple threads. </para> <section xml:id="tools.tokenizer.subword.sentencepiece"> <title>SentencePiece</title> 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 c6ad75542..0a08c980e 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 @@ -16,6 +16,7 @@ */ package opennlp.subword.sentencepiece; +import java.io.Serializable; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.HashMap; @@ -31,7 +32,9 @@ import java.util.PriorityQueue; * <p>Only pieces of the normal, user-defined, and unused types participate in merges; a merge that * lands on an unused piece is re-segmented back into its constituents.</p> */ -final class BpeEncoder { +final class BpeEncoder implements Serializable { + + private static final long serialVersionUID = -57799941356582785L; private static final int MAX_RESEGMENT_DEPTH = 100; @@ -77,8 +80,12 @@ final class BpeEncoder { * @param normalized The buffer holding the normalized UTF-8 bytes; must not be null. * @param size The number of valid bytes in {@code normalized}. * @return The segments covering all bytes, in text order. + * @throws IllegalArgumentException Thrown if {@code normalized} is null. */ List<Segment> encode(byte[] normalized, int size) { + if (normalized == null) { + throw new IllegalArgumentException("The normalized buffer must not be null."); + } if (size == 0) { return List.of(); } @@ -92,7 +99,7 @@ final class BpeEncoder { while (position < size) { int matched = 0; if (userDefinedMatcher != null) { - matched = longestUserDefinedMatch(normalized, size, position); + matched = userDefinedMatcher.longestMatch(normalized, size, position); } final boolean frozen = matched > 0; final int length = frozen ? matched @@ -158,6 +165,20 @@ final class BpeEncoder { return output; } + /** + * Offers the adjacent symbol pair {@code (left, right)} as a merge candidate: the pair joins + * the agenda only when the concatenation is a mergeable vocabulary piece, and a merge landing + * on an unused piece is remembered in {@code revMerge} for later re-segmentation. + * + * @param normalized The buffer holding the normalized UTF-8 bytes. + * @param from Per symbol, the inclusive start offset in {@code normalized}. + * @param to Per symbol, the exclusive end offset in {@code normalized}. + * @param freeze Per symbol, whether it is a user-defined symbol excluded from merging. + * @param left The index of the left symbol, or {@code -1} for none. + * @param right The index of the right symbol, or {@code -1} for none. + * @param agenda The merge agenda to add to. + * @param revMerge The map from a merged piece to its two constituents. + */ private void maybeAddPair(byte[] normalized, int[] from, int[] to, boolean[] freeze, int left, int right, PriorityQueue<Pair> agenda, Map<String, String[]> revMerge) { @@ -207,19 +228,4 @@ final class BpeEncoder { consumed = resegment(parts[0], consumed, depth + 1, revMerge, output); return resegment(parts[1], consumed, depth + 1, revMerge, output); } - - private int longestUserDefinedMatch(byte[] input, int inputLength, int from) { - int node = userDefinedMatcher.root(); - int longest = 0; - for (int i = from; i < inputLength; i++) { - node = userDefinedMatcher.step(node, input[i]); - if (node == PieceTrie.DEAD) { - break; - } - if (userDefinedMatcher.value(node) >= 0) { - longest = i - from + 1; - } - } - return longest; - } } diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/ByteBuilder.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/ByteBuilder.java index 95a02fddd..ef08bb624 100644 --- a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/ByteBuilder.java +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/ByteBuilder.java @@ -68,9 +68,15 @@ final class ByteBuilder { /** * Shrinks the valid length. * - * @param newLength The new length, not greater than the current length. + * @param newLength The new length, not negative and not greater than the current length. + * @throws IllegalArgumentException Thrown if {@code newLength} is negative or greater than the + * current length. */ void truncate(int newLength) { + if (newLength < 0 || newLength > length) { + throw new IllegalArgumentException( + "The new length " + newLength + " is outside [0, " + length + "]."); + } length = newLength; } 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 0ffb5e8be..28eefd3f1 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 @@ -16,6 +16,8 @@ */ package opennlp.subword.sentencepiece; +import java.io.Serializable; + /** * Read-only lookup over a serialized Darts-clone double-array trie, the dictionary format * embedded in a SentencePiece model's precompiled character map. @@ -25,8 +27,23 @@ package opennlp.subword.sentencepiece; * prefix match is needed here, so this walks the byte key once and remembers the last accepting * state. Out-of-range unit references, which a well-formed trie never produces, fail loudly * rather than reading arbitrary memory.</p> + * + * @see <a href="https://github.com/s-yata/darts-clone";>Darts-clone</a> */ -final class DoubleArrayTrie { +final class DoubleArrayTrie implements Serializable { + + private static final long serialVersionUID = -1572336116472261588L; + + // A non-leaf unit stores its transition label in the low 8 bits and the leaf flag in the sign + // bit. Key bytes are in [0, 255] with the sign bit clear, so comparing (unit & this mask) + // against a key byte both matches the label and rejects leaf units in one test. + private static final int LEAF_FLAG_AND_LABEL_MASK = 0x800000FF; + + // A leaf unit stores the key's value in its low 31 bits; the sign bit is the leaf flag. + private static final int LEAF_VALUE_MASK = 0x7FFFFFFF; + + // Bit 8 of a non-leaf unit marks that one of its children is a leaf holding this key's value. + private static final int HAS_LEAF_BIT = 8; private final int[] units; @@ -72,12 +89,12 @@ final class DoubleArrayTrie { final int b = key[i] & 0xFF; nodePos ^= b; unit = u[nodePos]; - if ((unit & 0x800000FF) != b) { + if ((unit & LEAF_FLAG_AND_LABEL_MASK) != b) { return result; } nodePos ^= offset(unit); - if (((unit >>> 8) & 1) == 1) { - final int value = u[nodePos] & 0x7FFFFFFF; + if (((unit >>> HAS_LEAF_BIT) & 1) == 1) { + final int value = u[nodePos] & LEAF_VALUE_MASK; result = ((long) value << 32) | (i - from + 1); } } @@ -101,11 +118,14 @@ final class DoubleArrayTrie { if (nodePos < 0 || nodePos >= units.length) { return false; } - return (units[nodePos] & 0x800000FF) == b; + return (units[nodePos] & LEAF_FLAG_AND_LABEL_MASK) == b; } /** - * Returns the offset from a unit to its children, as encoded by Darts-clone. + * Returns the offset from a unit to its children, as encoded by Darts-clone: bits 10 to 30 hold + * the raw offset, and bit 9 is an extension flag that scales it by 256 for far-away children. + * The expression {@code (unit & (1 << 9)) >>> 6} evaluates to 8 exactly when bit 9 is set, so + * the raw offset is shifted left by either 0 or 8 bits. * * @param unit The unit word. * @return The child offset. diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/IntBuilder.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/IntBuilder.java index 25dad2ec9..2da34c29b 100644 --- a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/IntBuilder.java +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/IntBuilder.java @@ -67,9 +67,15 @@ final class IntBuilder { /** * Shrinks the valid length. * - * @param newLength The new length, not greater than the current length. + * @param newLength The new length, not negative and not greater than the current length. + * @throws IllegalArgumentException Thrown if {@code newLength} is negative or greater than the + * current length. */ void truncate(int newLength) { + if (newLength < 0 || newLength > length) { + throw new IllegalArgumentException( + "The new length " + newLength + " is outside [0, " + length + "]."); + } length = newLength; } 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 4122746cb..c4c2a1472 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 @@ -28,6 +28,10 @@ import java.util.List; * 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.</p> + * + * @see <a href= + * "https://github.com/google/sentencepiece/blob/master/src/sentencepiece_model.proto";> + * sentencepiece_model.proto</a> */ final class ModelProtoReader { @@ -37,9 +41,43 @@ final class ModelProtoReader { private static final int WIRE_LEN = 2; private static final int WIRE_FIXED32 = 5; + // Field numbers of the ModelProto message in sentencepiece_model.proto. + private static final int FIELD_MODEL_PIECES = 1; + private static final int FIELD_MODEL_TRAINER_SPEC = 2; + private static final int FIELD_MODEL_NORMALIZER_SPEC = 3; + private static final int FIELD_MODEL_SELF_TEST_DATA = 4; + + // Field numbers of the ModelProto.SentencePiece sub-message. + private static final int FIELD_PIECE_PIECE = 1; + private static final int FIELD_PIECE_SCORE = 2; + private static final int FIELD_PIECE_TYPE = 3; + + // Field numbers of the TrainerSpec sub-message. + private static final int FIELD_TRAINER_MODEL_TYPE = 3; + private static final int FIELD_TRAINER_TREAT_WHITESPACE_AS_SUFFIX = 24; + private static final int FIELD_TRAINER_BYTE_FALLBACK = 35; + private static final int FIELD_TRAINER_UNK_ID = 40; + + // Field numbers of the NormalizerSpec sub-message. + private static final int FIELD_NORMALIZER_PRECOMPILED_CHARSMAP = 2; + private static final int FIELD_NORMALIZER_ADD_DUMMY_PREFIX = 3; + private static final int FIELD_NORMALIZER_REMOVE_EXTRA_WHITESPACES = 4; + private static final int FIELD_NORMALIZER_ESCAPE_WHITESPACES = 5; + + // Field numbers of the SelfTestData sub-message and its Sample entries. + private static final int FIELD_SELF_TEST_SAMPLES = 1; + private static final int FIELD_SAMPLE_INPUT = 1; + private static final int FIELD_SAMPLE_EXPECTED = 2; + private final byte[] data; private int pos; + /** + * Prepares a reader positioned at the start of the given bytes; {@link #read(byte[])} drives + * the actual parse. + * + * @param data The raw bytes of a {@code .model} file. + */ private ModelProtoReader(byte[] data) { this.data = data; } @@ -61,10 +99,10 @@ final class ModelProtoReader { final long tag = reader.varint(); final int field = (int) (tag >>> 3); switch (field) { - case 1 -> reader.piece(model, reader.lenPayload(tag)); - case 2 -> reader.trainerSpec(model, reader.lenPayload(tag)); - case 3 -> reader.normalizerSpec(model, reader.lenPayload(tag)); - case 4 -> reader.selfTestData(model, reader.lenPayload(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); } } @@ -87,9 +125,9 @@ final class ModelProtoReader { while (pos < end) { final long tag = varint(); switch ((int) (tag >>> 3)) { - case 1 -> piece = utf8(lenPayload(tag)); - case 2 -> score = fixed32Float(tag); - case 3 -> type = (int) varintOf(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); } } @@ -115,10 +153,11 @@ final class ModelProtoReader { while (pos < end) { final long tag = varint(); switch ((int) (tag >>> 3)) { - case 3 -> model.modelType = (int) varintOf(tag); - case 24 -> model.treatWhitespaceAsSuffix = varintOf(tag) != 0; - case 35 -> model.byteFallback = varintOf(tag) != 0; - case 40 -> model.unkId = (int) varintOf(tag); + case FIELD_TRAINER_MODEL_TYPE -> model.modelType = (int) varintOf(tag); + case FIELD_TRAINER_TREAT_WHITESPACE_AS_SUFFIX -> + model.treatWhitespaceAsSuffix = varintOf(tag) != 0; + case FIELD_TRAINER_BYTE_FALLBACK -> model.byteFallback = varintOf(tag) != 0; + case FIELD_TRAINER_UNK_ID -> model.unkId = (int) varintOf(tag); default -> skip(tag); } } @@ -135,10 +174,12 @@ final class ModelProtoReader { while (pos < end) { final long tag = varint(); switch ((int) (tag >>> 3)) { - case 2 -> model.precompiledCharsMap = bytes(lenPayload(tag)); - case 3 -> model.addDummyPrefix = varintOf(tag) != 0; - case 4 -> model.removeExtraWhitespaces = varintOf(tag) != 0; - case 5 -> model.escapeWhitespaces = varintOf(tag) != 0; + case FIELD_NORMALIZER_PRECOMPILED_CHARSMAP -> + model.precompiledCharsMap = bytes(lenPayload(tag)); + case FIELD_NORMALIZER_ADD_DUMMY_PREFIX -> model.addDummyPrefix = varintOf(tag) != 0; + case FIELD_NORMALIZER_REMOVE_EXTRA_WHITESPACES -> + model.removeExtraWhitespaces = varintOf(tag) != 0; + case FIELD_NORMALIZER_ESCAPE_WHITESPACES -> model.escapeWhitespaces = varintOf(tag) != 0; default -> skip(tag); } } @@ -154,15 +195,15 @@ final class ModelProtoReader { private void selfTestData(RawModel model, int end) { while (pos < end) { final long tag = varint(); - if ((int) (tag >>> 3) == 1) { + if ((int) (tag >>> 3) == FIELD_SELF_TEST_SAMPLES) { final int sampleEnd = lenPayload(tag); String input = null; String expected = null; while (pos < sampleEnd) { final long sampleTag = varint(); switch ((int) (sampleTag >>> 3)) { - case 1 -> input = utf8(lenPayload(sampleTag)); - case 2 -> expected = utf8(lenPayload(sampleTag)); + case FIELD_SAMPLE_INPUT -> input = utf8(lenPayload(sampleTag)); + case FIELD_SAMPLE_EXPECTED -> expected = utf8(lenPayload(sampleTag)); default -> skip(sampleTag); } } @@ -196,6 +237,13 @@ final class ModelProtoReader { return pos + (int) length; } + /** + * Reads the varint value of a field after checking its wire type. + * + * @param tag The field tag, whose wire type must be varint. + * @return The decoded value. + * @throws IllegalArgumentException Thrown if the wire type is wrong or the varint is malformed. + */ private long varintOf(long tag) { if ((tag & 7) != WIRE_VARINT) { throw malformed("field " + (tag >>> 3) + " is not a varint"); @@ -203,6 +251,13 @@ final class ModelProtoReader { return varint(); } + /** + * Reads the little-endian 32-bit float value of a field after checking its wire type. + * + * @param tag The field tag, whose wire type must be 32-bit. + * @return The decoded float. + * @throws IllegalArgumentException Thrown if the wire type is wrong or the input is truncated. + */ private float fixed32Float(long tag) { if ((tag & 7) != WIRE_FIXED32) { throw malformed("field " + (tag >>> 3) + " is not a 32-bit value"); @@ -216,12 +271,24 @@ final class ModelProtoReader { return Float.intBitsToFloat(bits); } + /** + * Decodes the bytes from the current position up to {@code end} as UTF-8, advancing past them. + * + * @param end The exclusive end offset of the payload. + * @return The decoded string. + */ private String utf8(int end) { final String s = new String(data, pos, end - pos, StandardCharsets.UTF_8); pos = end; return s; } + /** + * Copies the bytes from the current position up to {@code end}, advancing past them. + * + * @param end The exclusive end offset of the payload. + * @return The copied bytes. + */ private byte[] bytes(int end) { final byte[] b = new byte[end - pos]; System.arraycopy(data, pos, b, 0, b.length); @@ -268,6 +335,12 @@ final class ModelProtoReader { } } + /** + * Advances the position by a fixed number of bytes. + * + * @param count The number of bytes to skip. + * @throws IllegalArgumentException Thrown if fewer than {@code count} bytes remain. + */ private void advance(int count) { if (pos + count > data.length) { throw malformed("truncated field"); @@ -275,6 +348,12 @@ final class ModelProtoReader { pos += count; } + /** + * Creates the exception for malformed input, carrying the current byte position. + * + * @param detail A short description of what is malformed. + * @return The exception to throw. + */ private IllegalArgumentException malformed(String detail) { return new IllegalArgumentException( "The model data is malformed at byte " + pos + ": " + detail + "."); 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 795a251bc..5253a9467 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 @@ -16,6 +16,8 @@ */ package opennlp.subword.sentencepiece; +import java.io.Serializable; +import java.nio.charset.StandardCharsets; import java.util.Arrays; import java.util.Comparator; @@ -27,7 +29,9 @@ import java.util.Comparator; * 256-entry direct table and narrow nodes scan a short sorted label slice; both layouts enumerate * identical transitions.</p> */ -final class PieceTrie { +final class PieceTrie implements Serializable { + + private static final long serialVersionUID = 30340094783102906L; /** The node id returned when no transition exists. */ static final int DEAD = -1; @@ -45,6 +49,15 @@ final class PieceTrie { private final int[] directStart; private final int[] directPool; + /** + * Wraps the packed arrays produced by {@link Builder} and derives the direct-dispatch tables + * for wide nodes. + * + * @param childStart Per node, the start of its edge slice; one trailing entry marks the end. + * @param labels The transition label of every edge. + * @param childNodes The target node of every edge, parallel to {@code labels}. + * @param values Per node, the accepted piece id, or {@code -1}. + */ private PieceTrie(int[] childStart, byte[] labels, int[] childNodes, int[] values) { this.childStart = childStart; this.labels = labels; @@ -61,7 +74,7 @@ final class PieceTrie { } } this.directPool = new int[wide * 256]; - java.util.Arrays.fill(directPool, DEAD); + Arrays.fill(directPool, DEAD); for (int node = 0; node < values.length; node++) { final int direct = directStart[node]; if (direct >= 0) { @@ -131,6 +144,41 @@ final class PieceTrie { return values[node]; } + /** + * Returns the byte length of the longest piece in this trie that is a prefix of + * {@code input[from, inputLength)}. + * + * @param input The UTF-8 input buffer; must not be null. + * @param inputLength The number of valid bytes in {@code input}. + * @param from The offset to match from. + * @return The matched length in bytes, or zero when no piece matches. + */ + int longestMatch(byte[] input, int inputLength, int from) { + int node = root(); + int longest = 0; + for (int i = from; i < inputLength; i++) { + node = step(node, input[i]); + if (node == DEAD) { + break; + } + if (value(node) >= 0) { + longest = i - from + 1; + } + } + return longest; + } + + /** + * Creates the exception reported wherever a vocabulary piece turns out to be defined twice, + * keeping the message identical across all detection sites. + * + * @param piece The duplicated piece content. + * @return The exception to throw. + */ + static IllegalArgumentException duplicatePiece(String piece) { + return new IllegalArgumentException("The piece '" + piece + "' is defined more than once."); + } + // Builds the packed form from keys sorted by unsigned byte order. Key ranges sharing a prefix // are contiguous after the sort, so each recursion partitions its range by the byte at the // current depth. @@ -150,6 +198,14 @@ final class PieceTrie { private int nextNode; private int nextEdge; + /** + * Prepares a builder over the keys and their sort order; {@link #count} and {@link #fill} + * perform the actual construction. + * + * @param pieces The UTF-8 bytes of each piece. + * @param ids The id stored for each piece, parallel to {@code pieces}. + * @param order The indices of {@code pieces} sorted by unsigned byte order. + */ Builder(byte[][] pieces, int[] ids, Integer[] order) { this.pieces = pieces; this.ids = ids; @@ -172,9 +228,7 @@ final class PieceTrie { i++; // A second key ending at the same depth is a duplicate; the sort made them adjacent. if (i < to && pieces[order[i]].length == depth) { - throw new IllegalArgumentException("The piece '" - + new String(pieces[order[i]], java.nio.charset.StandardCharsets.UTF_8) - + "' is defined more than once."); + throw duplicatePiece(new String(pieces[order[i]], StandardCharsets.UTF_8)); } } while (i < to) { @@ -213,9 +267,7 @@ final class PieceTrie { int i = from; if (i < to && pieces[order[i]].length == depth) { if (values[node] != -1 || (i + 1 < to && pieces[order[i + 1]].length == depth)) { - throw new IllegalArgumentException( - "The piece '" + new String(pieces[order[i]], java.nio.charset.StandardCharsets.UTF_8) - + "' is defined more than once."); + throw duplicatePiece(new String(pieces[order[i]], StandardCharsets.UTF_8)); } values[node] = ids[order[i]]; i++; diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/SentencePieceNormalizer.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/SentencePieceNormalizer.java index 4d5589514..0dd83c278 100644 --- a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/SentencePieceNormalizer.java +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/SentencePieceNormalizer.java @@ -16,6 +16,8 @@ */ package opennlp.subword.sentencepiece; +import java.io.Serializable; + /** * The model-embedded text normalizer of a SentencePiece model, operating in UTF-8 byte space. * @@ -26,7 +28,9 @@ package opennlp.subword.sentencepiece; * was derived from, with one trailing entry for the end position; that map is what lets every * downstream piece report an exact span of the caller's text.</p> */ -final class SentencePieceNormalizer { +final class SentencePieceNormalizer implements Serializable { + + private static final long serialVersionUID = -3059745470932191300L; // U+2581 LOWER ONE EIGHTH BLOCK in UTF-8, the escaped form of a space. static final byte[] SPACE_SYMBOL = {(byte) 0xE2, (byte) 0x96, (byte) 0x81}; @@ -238,6 +242,15 @@ final class SentencePieceNormalizer { private static final byte[] SINGLE_SPACE = {' '}; + /** + * Appends the space symbol to the normalized output, mapping each of its bytes to the same + * original-byte offset. + * + * @param normalized The normalized-byte builder to append to. + * @param normToOrig The offset-map builder to append to. + * @param spaceSymbol The bytes of the (possibly escaped) space symbol. + * @param consumed The original-byte offset the symbol maps back to. + */ private static void appendSpace(ByteBuilder normalized, IntBuilder normToOrig, byte[] spaceSymbol, int consumed) { normalized.append(spaceSymbol, 0, spaceSymbol.length); @@ -259,7 +272,7 @@ final class SentencePieceNormalizer { */ private void normalizePrefix(byte[] input, int inputLength, int from, Chunk chunk) { if (userDefinedMatcher != null) { - final int matched = longestUserDefinedMatch(input, inputLength, from); + final int matched = userDefinedMatcher.longestMatch(input, inputLength, from); if (matched > 0) { chunk.data = input; chunk.from = from; @@ -303,30 +316,6 @@ final class SentencePieceNormalizer { chunk.consumed = charLength; } - /** - * Returns the byte length of the longest user-defined symbol that is a prefix of - * {@code input[from, inputLength)}. - * - * @param input The UTF-8 input buffer. - * @param inputLength The number of valid bytes in {@code input}. - * @param from The offset to match from. - * @return The matched length in bytes, or zero when no user-defined symbol matches. - */ - private int longestUserDefinedMatch(byte[] input, int inputLength, int from) { - int node = userDefinedMatcher.root(); - int longest = 0; - for (int i = from; i < inputLength; i++) { - node = userDefinedMatcher.step(node, input[i]); - if (node == PieceTrie.DEAD) { - break; - } - if (userDefinedMatcher.value(node) >= 0) { - longest = i - from + 1; - } - } - return longest; - } - /** * Returns the byte length of a UTF-8 sequence from its lead byte; trail and malformed lead bytes * report one byte. 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 961bb58b3..e8c0e04ee 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 @@ -22,9 +22,11 @@ import java.nio.charset.StandardCharsets; import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; +import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.IntUnaryOperator; import opennlp.tools.tokenize.SubwordPiece; import opennlp.tools.tokenize.SubwordTokenizer; @@ -43,6 +45,11 @@ import opennlp.tools.util.normalizer.OffsetAwareNormalizer; * for reuse outside tokenization.</p> * * <p>Instances are immutable after loading and safe for concurrent use by multiple threads.</p> + * + * @see <a href="https://github.com/google/sentencepiece";>SentencePiece</a> + * @see <a href="https://aclanthology.org/D18-2012/";>Kudo & Richardson (EMNLP 2018), + * "SentencePiece: A simple and language independent subword tokenizer and detokenizer for + * Neural Text Processing"</a> */ public final class SentencePieceTokenizer implements SubwordTokenizer, OffsetAwareNormalizer { @@ -82,6 +89,13 @@ public final class SentencePieceTokenizer implements SubwordTokenizer, OffsetAwa private final List<String> selfTestInputs; private final List<String> selfTestExpected; + /** + * Validates a parsed model and derives the runtime structures: the piece maps, the byte-piece + * table, the normalizer, and the encoder matching the model's algorithm. + * + * @param model The parsed model description. + * @throws IllegalArgumentException Thrown if the model is structurally invalid. + */ private SentencePieceTokenizer(ModelProtoReader.RawModel model) { final int count = model.pieces.size(); pieces = model.pieces.toArray(new String[0]); @@ -106,7 +120,7 @@ public final class SentencePieceTokenizer implements SubwordTokenizer, OffsetAwa reservedPieces = new HashMap<>(); final List<String> userDefined = new ArrayList<>(); byteToId = new int[256]; - java.util.Arrays.fill(byteToId, -1); + Arrays.fill(byteToId, -1); int foundUnkId = -1; float minScore = Float.MAX_VALUE; for (int i = 0; i < count; i++) { @@ -124,7 +138,7 @@ public final class SentencePieceTokenizer implements SubwordTokenizer, OffsetAwa final Map<String, Integer> target = isMain || algorithm == Algorithm.BPE ? mainPieces : reservedPieces; if (mainPieces.containsKey(piece) || reservedPieces.containsKey(piece)) { - throw new IllegalArgumentException("The piece '" + piece + "' is defined more than once."); + throw PieceTrie.duplicatePiece(piece); } target.put(piece, i); switch (types[i]) { @@ -165,7 +179,8 @@ public final class SentencePieceTokenizer implements SubwordTokenizer, OffsetAwa } } - final PieceTrie userDefinedMatcher = userDefined.isEmpty() ? null : trieOf(userDefined, id -> 0); + final PieceTrie userDefinedMatcher = + userDefined.isEmpty() ? null : trieOf(userDefined, id -> 0); normalizer = new SentencePieceNormalizer(model.precompiledCharsMap, model.addDummyPrefix, model.removeExtraWhitespaces, model.escapeWhitespaces, model.treatWhitespaceAsSuffix, @@ -204,8 +219,14 @@ public final class SentencePieceTokenizer implements SubwordTokenizer, OffsetAwa selfTestExpected = List.copyOf(model.selfTestExpected); } - private static PieceTrie trieOf(List<String> pieceList, - java.util.function.IntUnaryOperator idOf) { + /** + * Builds a {@link PieceTrie} over the given pieces. + * + * @param pieceList The pieces to index. + * @param idOf Maps a piece's index in {@code pieceList} to the id the trie stores for it. + * @return The packed trie. + */ + private static PieceTrie trieOf(List<String> pieceList, IntUnaryOperator idOf) { final byte[][] keys = new byte[pieceList.size()][]; final int[] ids = new int[pieceList.size()]; for (int i = 0; i < keys.length; i++) { @@ -246,11 +267,7 @@ public final class SentencePieceTokenizer implements SubwordTokenizer, OffsetAwa return new SentencePieceTokenizer(ModelProtoReader.read(in.readAllBytes())); } - /** - * {@inheritDoc} - * - * @throws IllegalArgumentException Thrown if {@code text} is null. - */ + /** {@inheritDoc} */ @Override public List<SubwordPiece> encode(CharSequence text) { if (text == null) { @@ -526,13 +543,16 @@ public final class SentencePieceTokenizer implements SubwordTokenizer, OffsetAwa return selfTestExpected; } + // The prefix of a byte-fallback piece string; a full piece has the form "<0xAB>". + private static final String BYTE_PIECE_PREFIX = "<0x"; + // "<0xAB>" piece strings for all byte values, as byte fallback emits them. private static final String[] BYTE_PIECES = new String[256]; static { final char[] hex = "0123456789ABCDEF".toCharArray(); for (int b = 0; b < 256; b++) { - BYTE_PIECES[b] = "<0x" + hex[b >>> 4] + hex[b & 0xF] + ">"; + BYTE_PIECES[b] = BYTE_PIECE_PREFIX + hex[b >>> 4] + hex[b & 0xF] + ">"; } } @@ -543,7 +563,7 @@ public final class SentencePieceTokenizer implements SubwordTokenizer, OffsetAwa * @return The byte value in {@code [0, 255]}, or {@code -1} when the string is not a byte piece. */ private static int parseBytePiece(String piece) { - if (piece.length() != 6 || !piece.startsWith("<0x") || piece.charAt(5) != '>') { + if (piece.length() != 6 || !piece.startsWith(BYTE_PIECE_PREFIX) || piece.charAt(5) != '>') { return -1; } final int high = Character.digit(piece.charAt(3), 16); diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/UnigramEncoder.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/UnigramEncoder.java index c87662bef..1efa00310 100644 --- a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/UnigramEncoder.java +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/UnigramEncoder.java @@ -16,6 +16,7 @@ */ package opennlp.subword.sentencepiece; +import java.io.Serializable; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -27,9 +28,14 @@ import java.util.List; * <p>Characters no piece covers fall back to the unknown id with a fixed penalty below the lowest * piece score, and user-defined symbols receive a length-based bonus score so they always win.</p> */ -final class UnigramEncoder { +final class UnigramEncoder implements Serializable { + + private static final long serialVersionUID = 5648005733414803707L; private static final float UNK_PENALTY = 10.0f; + // The score of a user-defined symbol is this bonus per matched byte beyond the first instead + // of a trained log-probability, so longer user-defined matches always win the best path. + private static final float USER_DEFINED_LENGTH_BONUS = 0.1f; private static final float SCORE_RESET_THRESHOLD = 100000.0f; private final PieceTrie trie; @@ -65,8 +71,12 @@ final class UnigramEncoder { * @param normalized The buffer holding the normalized UTF-8 bytes; must not be null. * @param size The number of valid bytes in {@code normalized}. * @return The best-path segments covering all bytes, in text order. + * @throws IllegalArgumentException Thrown if {@code normalized} is null. */ List<Segment> encode(byte[] normalized, int size) { + if (normalized == null) { + throw new IllegalArgumentException("The normalized buffer must not be null."); + } if (size == 0) { return List.of(); } @@ -118,7 +128,8 @@ final class UnigramEncoder { maxFrontier = Math.max(maxFrontier, keyPos); final int length = keyPos - startsAt; // User-defined symbols receive a length bonus instead of a trained score. - final float score = userDefined[id] ? 0.1f * (length - 1) : scores[id]; + final float score = userDefined[id] + ? USER_DEFINED_LENGTH_BONUS * (length - 1) : scores[id]; final float candidate = score + bestScoreTillHere; final int slot = 3 * keyPos; if (best[slot] == -1 || candidate > Float.intBitsToFloat(best[slot + 1])) { diff --git a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/Utf8Text.java b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/Utf8Text.java index e830816e6..4f2bd95f0 100644 --- a/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/Utf8Text.java +++ b/opennlp-extensions/opennlp-subword/src/main/java/opennlp/subword/sentencepiece/Utf8Text.java @@ -32,6 +32,14 @@ final class Utf8Text { private final int[] byteToChar; private final int charLength; + /** + * Wraps an encoded buffer and its offset map. + * + * @param bytes The UTF-8 buffer. + * @param byteLength The number of valid bytes in {@code bytes}. + * @param byteToChar The byte-to-UTF-16 offset map, or null for pure-ASCII text. + * @param charLength The length of the original text in UTF-16 units. + */ private Utf8Text(byte[] bytes, int byteLength, int[] byteToChar, int charLength) { this.bytes = bytes; this.byteLength = byteLength; diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceFixtures.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceFixtures.java new file mode 100644 index 000000000..11e54d2be --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceFixtures.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.subword.sentencepiece; + +import java.io.BufferedReader; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +import opennlp.tools.tokenize.SubwordPiece; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** + * Shared support for the tab-separated parity fixture files produced by the + * {@code gen_fixtures.py} and {@code gen_real_fixtures.py} scripts in the test resources: one + * line per input, holding the input, the expected piece count, four columns per expected piece + * (content, id, start, end), and the expected normalized form. + */ +final class SentencePieceFixtures { + + private SentencePieceFixtures() { + } + + /** + * One parsed fixture line: an input with the piece sequence and normalized form the reference + * implementation produced for it. + * + * @param input The text to encode. + * @param pieces The expected pieces with ids and original-text spans, in text order. + * @param normalized The expected normalized form of {@code input}. + */ + record Fixture(String input, List<SubwordPiece> pieces, String normalized) { + } + + /** + * Reads all fixture lines from a reader. + * + * @param reader The reader positioned at the start of a fixture file; must not be null. + * @return The parsed fixtures in file order. + * @throws IOException Thrown if the reader fails. + */ + static List<Fixture> read(BufferedReader reader) throws IOException { + final List<Fixture> fixtures = new ArrayList<>(); + String line; + while ((line = reader.readLine()) != null) { + final String[] cols = line.split("\t", -1); + final String input = unescape(cols[0]); + final int count = Integer.parseInt(cols[1]); + final List<SubwordPiece> pieces = new ArrayList<>(count); + for (int i = 0; i < count; i++) { + pieces.add(new SubwordPiece(unescape(cols[2 + i * 4]), + Integer.parseInt(cols[3 + i * 4]), Integer.parseInt(cols[4 + i * 4]), + Integer.parseInt(cols[5 + i * 4]))); + } + fixtures.add(new Fixture(input, pieces, unescape(cols[2 + count * 4]))); + } + return fixtures; + } + + /** + * Asserts that a tokenizer reproduces one fixture exactly: the piece sequence with ids and + * spans, and the normalized form. + * + * @param tokenizer The tokenizer under test. + * @param fixture The expected encoding. + * @param context A prefix for failure messages that identifies the model and input. + */ + static void assertFixture(SentencePieceTokenizer tokenizer, Fixture fixture, String context) { + final List<SubwordPiece> actual = tokenizer.encode(fixture.input()); + assertEquals(fixture.pieces().size(), actual.size(), + context + " piece count; got " + actual); + for (int i = 0; i < actual.size(); i++) { + final SubwordPiece expected = fixture.pieces().get(i); + final SubwordPiece got = actual.get(i); + assertEquals(expected.piece(), got.piece(), context + " piece " + i); + assertEquals(expected.id(), got.id(), context + " id of piece " + i); + assertEquals(expected.start(), got.start(), context + " start of piece " + i); + assertEquals(expected.end(), got.end(), context + " end of piece " + i); + } + assertEquals(fixture.normalized(), tokenizer.normalize(fixture.input()).toString(), + context + " normalized form"); + } + + /** + * Reverses the fixture files' escaping of tab, newline, carriage return, and backslash. + * + * @param s The escaped column content. + * @return The unescaped text. + * @throws IllegalArgumentException Thrown if an unknown escape sequence occurs. + */ + static String unescape(String s) { + final StringBuilder out = new StringBuilder(s.length()); + for (int i = 0; i < s.length(); i++) { + final char c = s.charAt(i); + if (c == '\\' && i + 1 < s.length()) { + i++; + switch (s.charAt(i)) { + case 't' -> out.append('\t'); + case 'n' -> out.append('\n'); + case 'r' -> out.append('\r'); + case '\\' -> out.append('\\'); + default -> throw new IllegalArgumentException("bad escape in fixture: " + s); + } + } else { + out.append(c); + } + } + return out.toString(); + } +} diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceModelValidationTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceModelValidationTest.java index 66595eff8..a65df0899 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 @@ -21,7 +21,9 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; import java.nio.charset.StandardCharsets; +import java.nio.file.Path; import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import java.util.concurrent.Callable; import java.util.concurrent.ExecutorService; @@ -45,7 +47,7 @@ class SentencePieceModelValidationTest { @Test void testNullAndEmptyInputFailLoudly() { assertThrows(IllegalArgumentException.class, - () -> SentencePieceTokenizer.load((java.nio.file.Path) null)); + () -> SentencePieceTokenizer.load((Path) null)); assertThrows(IllegalArgumentException.class, () -> SentencePieceTokenizer.load((InputStream) null)); assertThrows(IllegalArgumentException.class, @@ -62,7 +64,7 @@ class SentencePieceModelValidationTest { @Test void testTruncatedModelFailsLoudly() throws IOException { final byte[] whole = readModel(); - final byte[] truncated = java.util.Arrays.copyOf(whole, whole.length / 3); + final byte[] truncated = Arrays.copyOf(whole, whole.length / 3); assertThrows(IllegalArgumentException.class, () -> SentencePieceTokenizer.load(new ByteArrayInputStream(truncated))); } diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceParityTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceParityTest.java index 7d970e3c5..95730ef91 100644 --- a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceParityTest.java +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceParityTest.java @@ -21,7 +21,6 @@ import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; import java.nio.charset.StandardCharsets; -import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.StringJoiner; @@ -30,8 +29,6 @@ import java.util.concurrent.ConcurrentHashMap; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; -import opennlp.tools.tokenize.SubwordPiece; - import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -63,22 +60,10 @@ class SentencePieceParityTest { void testFixtureParity(String model) throws IOException { final SentencePieceTokenizer tokenizer = tokenizer(model); int lines = 0; - for (final Fixture fixture : fixtures(model)) { + for (final SentencePieceFixtures.Fixture fixture : fixtures(model)) { lines++; - final List<SubwordPiece> actual = tokenizer.encode(fixture.input); - final String context = model + " input <" + fixture.input + ">"; - assertEquals(fixture.pieces.size(), actual.size(), - context + " piece count; got " + actual); - for (int i = 0; i < actual.size(); i++) { - final SubwordPiece expected = fixture.pieces.get(i); - final SubwordPiece got = actual.get(i); - assertEquals(expected.piece(), got.piece(), context + " piece " + i); - assertEquals(expected.id(), got.id(), context + " id of piece " + i); - assertEquals(expected.start(), got.start(), context + " start of piece " + i); - assertEquals(expected.end(), got.end(), context + " end of piece " + i); - } - assertEquals(fixture.normalized, tokenizer.normalize(fixture.input).toString(), - context + " normalized form"); + SentencePieceFixtures.assertFixture(tokenizer, fixture, + model + " input <" + fixture.input() + ">"); } assertTrue(lines >= 30, "the fixture file must not be empty or truncated"); } @@ -101,50 +86,12 @@ class SentencePieceParityTest { } } - private record Fixture(String input, List<SubwordPiece> pieces, String normalized) { - } - - private static List<Fixture> fixtures(String model) throws IOException { - final List<Fixture> fixtures = new ArrayList<>(); + private static List<SentencePieceFixtures.Fixture> fixtures(String model) throws IOException { try (InputStream in = SentencePieceParityTest.class.getResourceAsStream(model + ".fixtures.tsv")) { assertNotNull(in, "missing test resource " + model + ".fixtures.tsv"); - final BufferedReader reader = - new BufferedReader(new InputStreamReader(in, StandardCharsets.UTF_8)); - String line; - while ((line = reader.readLine()) != null) { - final String[] cols = line.split("\t", -1); - final String input = unescape(cols[0]); - final int count = Integer.parseInt(cols[1]); - final List<SubwordPiece> pieces = new ArrayList<>(count); - for (int i = 0; i < count; i++) { - pieces.add(new SubwordPiece(unescape(cols[2 + i * 4]), - Integer.parseInt(cols[3 + i * 4]), Integer.parseInt(cols[4 + i * 4]), - Integer.parseInt(cols[5 + i * 4]))); - } - fixtures.add(new Fixture(input, pieces, unescape(cols[2 + count * 4]))); - } - } - return fixtures; - } - - private static String unescape(String s) { - final StringBuilder out = new StringBuilder(s.length()); - for (int i = 0; i < s.length(); i++) { - final char c = s.charAt(i); - if (c == '\\' && i + 1 < s.length()) { - i++; - switch (s.charAt(i)) { - case 't' -> out.append('\t'); - case 'n' -> out.append('\n'); - case 'r' -> out.append('\r'); - case '\\' -> out.append('\\'); - default -> throw new IllegalArgumentException("bad escape in fixture: " + s); - } - } else { - out.append(c); - } + return SentencePieceFixtures.read( + new BufferedReader(new InputStreamReader(in, StandardCharsets.UTF_8))); } - return out.toString(); } } diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceRealModelEvalTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceRealModelEvalTest.java index 1255827ba..fd1a9ca9e 100644 --- a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceRealModelEvalTest.java +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceRealModelEvalTest.java @@ -26,9 +26,6 @@ import java.util.stream.Stream; import org.junit.jupiter.api.Test; -import opennlp.tools.tokenize.SubwordPiece; - -import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assumptions.assumeTrue; @@ -64,49 +61,14 @@ class SentencePieceRealModelEvalTest { private static void assertModel(Path modelPath, Path fixturesPath) throws IOException { final SentencePieceTokenizer tokenizer = SentencePieceTokenizer.load(modelPath); - int lines = 0; + final List<SentencePieceFixtures.Fixture> fixtures; try (BufferedReader reader = Files.newBufferedReader(fixturesPath, StandardCharsets.UTF_8)) { - String line; - while ((line = reader.readLine()) != null) { - lines++; - final String[] cols = line.split("\t", -1); - final String input = unescape(cols[0]); - final int count = Integer.parseInt(cols[1]); - final String context = modelPath.getFileName() + " input <" + input + ">"; - - final List<SubwordPiece> actual = tokenizer.encode(input); - assertEquals(count, actual.size(), context + " piece count; got " + actual); - for (int i = 0; i < count; i++) { - final SubwordPiece got = actual.get(i); - assertEquals(unescape(cols[2 + i * 4]), got.piece(), context + " piece " + i); - assertEquals(Integer.parseInt(cols[3 + i * 4]), got.id(), context + " id " + i); - assertEquals(Integer.parseInt(cols[4 + i * 4]), got.start(), context + " start " + i); - assertEquals(Integer.parseInt(cols[5 + i * 4]), got.end(), context + " end " + i); - } - assertEquals(unescape(cols[2 + count * 4]), tokenizer.normalize(input).toString(), - context + " normalized form"); - } + fixtures = SentencePieceFixtures.read(reader); } - assertTrue(lines >= 30, modelPath.getFileName() + " fixtures must not be truncated"); - } - - private static String unescape(String s) { - final StringBuilder out = new StringBuilder(s.length()); - for (int i = 0; i < s.length(); i++) { - final char c = s.charAt(i); - if (c == '\\' && i + 1 < s.length()) { - i++; - switch (s.charAt(i)) { - case 't' -> out.append('\t'); - case 'n' -> out.append('\n'); - case 'r' -> out.append('\r'); - case '\\' -> out.append('\\'); - default -> throw new IllegalArgumentException("bad escape in fixture: " + s); - } - } else { - out.append(c); - } + for (final SentencePieceFixtures.Fixture fixture : fixtures) { + SentencePieceFixtures.assertFixture(tokenizer, fixture, + modelPath.getFileName() + " input <" + fixture.input() + ">"); } - return out.toString(); + assertTrue(fixtures.size() >= 30, modelPath.getFileName() + " fixtures must not be truncated"); } } diff --git a/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceTokenizerSerializationTest.java b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceTokenizerSerializationTest.java new file mode 100644 index 000000000..f1b56b769 --- /dev/null +++ b/opennlp-extensions/opennlp-subword/src/test/java/opennlp/subword/sentencepiece/SentencePieceTokenizerSerializationTest.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.subword.sentencepiece; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; + +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.assertIterableEquals; + +/** + * Asserts the {@code Serializable} contract inherited through + * {@code opennlp.tools.util.normalizer.CharSequenceNormalizer}: a tokenizer round-tripped + * through Java object serialization must encode and normalize exactly like the original. + */ +class SentencePieceTokenizerSerializationTest { + + private static final String[] INPUTS = { + "", + "The quick brown fox jumps over the lazy dog.", + " Hello world ", + "tokenization and segmentation", + "caf\u00e9 na\u00efve \u4e2d\u6587" + }; + + @ParameterizedTest + @ValueSource(strings = {"tiny-unigram", "tiny-unigram-bytefb", "tiny-bpe", + "tiny-unigram-identity", "tiny-unigram-suffix"}) + void testRoundTripPreservesEncoding(String model) throws IOException, ClassNotFoundException { + final SentencePieceTokenizer original = SentencePieceParityTest.tokenizer(model); + + final ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (ObjectOutputStream out = new ObjectOutputStream(bytes)) { + out.writeObject(original); + } + final SentencePieceTokenizer copy; + try (ObjectInputStream in = + new ObjectInputStream(new ByteArrayInputStream(bytes.toByteArray()))) { + copy = (SentencePieceTokenizer) in.readObject(); + } + + assertEquals(original.algorithm(), copy.algorithm(), model + " algorithm"); + assertEquals(original.vocabularySize(), copy.vocabularySize(), model + " vocabulary size"); + for (final String input : INPUTS) { + final String context = model + " input <" + input + ">"; + assertIterableEquals(original.encode(input), copy.encode(input), context + " pieces"); + assertEquals(original.normalize(input).toString(), copy.normalize(input).toString(), + context + " normalized form"); + } + } +}
