OPENNLP-1885: Pure-Java SentencePiece inference with exact original-text spans (opennlp-subword)

opennlp-api/src/main/java/opennlp/tools/tokenize/BertNormalization.java

@@ -19,7 +19,7 @@
 
 /**
  * Character classifications and text transforms of the reference BERT
- * {@code BasicTokenizer}, shared by {@link BertTokenizer} and
+ * {@code BasicTokenizer}, shared by {@link WordpieceEncoder} and
  * {@link WordpieceTokenizer}.
  */
 final class BertNormalization {

opennlp-api/src/main/java/opennlp/tools/tokenize/BertTokenizer.java

@@ -6,77 +6,47 @@
  * (the "License"); you may not use this file except in compliance with
  * the License. You may obtain a copy of the License at
  *
- *   http://www.apache.org/licenses/LICENSE-2.0
+ *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package opennlp.tools.tokenize;
 
-import java.text.Normalizer;
-import java.util.Locale;
-import java.util.Objects;
+import java.util.ArrayList;
 import java.util.Set;
 
 import opennlp.tools.util.Span;
 
 /**
  * A {@link Tokenizer} implementation of the full BERT tokenization pipeline:
- * basic tokenization (text normalization) followed by wordpiece tokenization.
- * <p>
- * The basic tokenization stage reproduces the reference BERT
- * {@code BasicTokenizer}:
- * <ol>
- *  <li>Removal of control characters and normalization of all whitespace
- *      to single spaces.</li>
- *  <li>Whitespace isolation of CJK ideographs.</li>
- *  <li>For <i>uncased</i> models: lower casing and accent stripping
- *      (Unicode NFD decomposition with removal of combining marks).</li>
- *  <li>Isolation of every punctuation character as its own token.</li>
- * </ol>
- * The normalized text is then split into subwords by a
- * {@link WordpieceTokenizer} sharing the same vocabulary and special tokens.
- * <p>
- * This pipeline is required for correct results with BERT-style models:
- * feeding raw text directly to {@link WordpieceTokenizer} maps every token
- * that does not literally appear in the vocabulary - for uncased models that
- * includes every capitalized word - to the unknown token.
- * <p>
- * Whether to use the lower casing variant is a property of the model: uncased
- * models (for example {@code bert-base-uncased} and the
- * {@code sentence-transformers} models derived from it) require it, cased
- * models must not use it. Accent stripping is coupled to lower casing, as in
- * the reference implementation's default ({@code strip_accents} follows
- * {@code do_lower_case} unless overridden).
- * <p>
- * For reference see:
- * <ul>
- *  <li><a href="https://github.com/google-research/bert">
- *    https://github.com/google-research/bert</a> ({@code tokenization.py})</li>
- * </ul>
+ * basic tokenization (text normalization) followed by wordpiece tokenization,
+ * with the classification and separator tokens framing every result.
+ *
+ * @deprecated Use {@link WordpieceEncoder} instead:
+ *     {@link WordpieceEncoder#encodeToPieces(CharSequence)} returns the same
+ *     {@code String[]} as {@link #tokenize(String)}, and {@code encode} additionally
+ *     carries vocabulary ids and original-text spans. This class is scheduled for
+ *     removal after one stable release.
  *
- * @see WordpieceTokenizer
+ * @see WordpieceEncoder
  */
+@Deprecated(since = "3.0.0", forRemoval = true)
 public class BertTokenizer implements Tokenizer {
 
-  /**
-   * Maximum characters per word before the word is replaced with the unknown
-   * token, matching the reference BERT implementation.
-   */
-  private static final int MAX_WORD_CHARACTERS = 100;
-
-  private final WordpieceTokenizer wordpieceTokenizer;
-  private final boolean lowerCase;
+  private final WordpieceEncoder encoder;
 
   /**
    * Initializes a {@link BertTokenizer} for an <i>uncased</i> 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<String> vocabulary) {
     this(vocabulary, true);
@@ -88,6 +58,9 @@ public BertTokenizer(Set<String> vocabulary) {
    * @param vocabulary The wordpiece vocabulary. Must not be {@code null}.
    * @param lowerCase  {@code true} for uncased models (lower casing and accent
    *                   stripping), {@code false} for cased models.
+   *
+   * @throws IllegalArgumentException Thrown if the vocabulary is {@code null},
+   *     contains {@code null}, or is missing a BERT special token.
    */
   public BertTokenizer(Set<String> vocabulary, boolean lowerCase) {
     this(vocabulary, lowerCase, WordpieceTokenizer.BERT_CLS_TOKEN,
@@ -101,19 +74,23 @@ public BertTokenizer(Set<String> vocabulary, boolean lowerCase) {
    * @param vocabulary          The wordpiece vocabulary. Must not be {@code null}.
    * @param lowerCase           {@code true} for uncased models (lower casing and
    *                            accent stripping), {@code false} for cased models.
-   * @param classificationToken The CLS token.
-   * @param separatorToken      The SEP token.
-   * @param unknownToken        The UNK token.
+   * @param classificationToken The CLS token; must be in the vocabulary.
+   * @param separatorToken      The SEP token; must be in the vocabulary.
+   * @param unknownToken        The UNK token; must be in the vocabulary.
+   *
+   * @throws IllegalArgumentException Thrown if any argument is {@code null},
+   *     the vocabulary contains {@code null}, or a special token is missing
+   *     from the vocabulary.
    */
   public BertTokenizer(Set<String> vocabulary, boolean lowerCase,
       String classificationToken, String separatorToken, String unknownToken) {
-    Objects.requireNonNull(vocabulary, "vocabulary must not be null");
-    Objects.requireNonNull(classificationToken, "classificationToken must not be null");
-    Objects.requireNonNull(separatorToken, "separatorToken must not be null");
-    Objects.requireNonNull(unknownToken, "unknownToken must not be null");
-    this.wordpieceTokenizer = new WordpieceTokenizer(vocabulary,
-        classificationToken, separatorToken, unknownToken, MAX_WORD_CHARACTERS);
-    this.lowerCase = lowerCase;
+    if (vocabulary == null) {
+      throw new IllegalArgumentException("vocabulary must not be null");
+    }
+    // The encoder assigns each piece its list index as the id. Ids are unused on the
+    // tokenize() path, so synthesizing them from an arbitrary set order is fine.
+    this.encoder = new WordpieceEncoder(new ArrayList<>(vocabulary), lowerCase,
+        classificationToken, separatorToken, unknownToken);
   }
 
   /**
@@ -123,15 +100,19 @@ public BertTokenizer(Set<String> vocabulary, boolean lowerCase,
    * @param text The text to tokenize. Must not be {@code null}.
    *
    * @return The wordpiece tokens.
+   *
+   * @throws IllegalArgumentException Thrown if {@code text} is {@code null}.
    */
   @Override
   public String[] tokenize(String text) {
-    return wordpieceTokenizer.tokenize(normalize(text));
+    return encoder.encodeToPieces(text);
   }
 
   /**
    * Not supported: wordpiece tokens (subwords, {@code ##} continuations and
    * special tokens) have no faithful character spans in the original text.
+   * Use {@link WordpieceEncoder#encode(CharSequence)} for pieces with
+   * original-text spans.
    *
    * @throws UnsupportedOperationException Always.
    */
@@ -141,71 +122,4 @@ public Span[] tokenizePos(String text) {
         "Wordpiece tokens cannot be mapped to character spans of the original text");
   }
 
-  /**
-   * Applies the BERT basic tokenization (normalization) stage.
-   *
-   * @param text The text to normalize. Must not be {@code null}.
-   *
-   * @return The normalized text, ready for wordpiece tokenization.
-   */
-  String normalize(String text) {
-    Objects.requireNonNull(text, "text must not be null");
-    String normalized = cleanText(text);
-    normalized = isolateCjkCharacters(normalized);
-    if (lowerCase) {
-      normalized = stripAccents(normalized.toLowerCase(Locale.ROOT));
-    }
-    return BertNormalization.isolatePunctuation(normalized);
-  }
-
-  /**
-   * Removes invalid and control characters and normalizes all whitespace
-   * characters to plain spaces.
-   */
-  private static String cleanText(String text) {
-    final StringBuilder cleaned = new StringBuilder(text.length());
-    text.codePoints().forEach(codePoint -> {
-      if (codePoint == 0 || codePoint == 0xFFFD || BertNormalization.isControl(codePoint)) {
-        return;
-      }
-      if (BertNormalization.isWhitespace(codePoint)) {
-        cleaned.append(' ');
-      } else {
-        cleaned.appendCodePoint(codePoint);
-      }
-    });
-    return cleaned.toString();
-  }
-
-  /**
-   * Surrounds every CJK ideograph with spaces, so each ideograph becomes its
-   * own token, matching the reference BERT treatment of Chinese text.
-   */
-  private static String isolateCjkCharacters(String text) {
-    final StringBuilder spaced = new StringBuilder(text.length());
-    text.codePoints().forEach(codePoint -> {
-      if (BertNormalization.isCjk(codePoint)) {
-        spaced.append(' ').appendCodePoint(codePoint).append(' ');
-      } else {
-        spaced.appendCodePoint(codePoint);
-      }
-    });
-    return spaced.toString();
-  }
-
-  /**
-   * Removes accents by Unicode NFD decomposition followed by removal of
-   * combining marks ({@code Mn}).
-   */
-  private static String stripAccents(String text) {
-    final String decomposed = Normalizer.normalize(text, Normalizer.Form.NFD);
-    final StringBuilder stripped = new StringBuilder(decomposed.length());
-    decomposed.codePoints().forEach(codePoint -> {
-      if (Character.getType(codePoint) != Character.NON_SPACING_MARK) {
-        stripped.appendCodePoint(codePoint);
-      }
-    });
-    return stripped.toString();
-  }
-
 }

opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordPiece.java

@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package opennlp.tools.tokenize;
+
+import opennlp.tools.util.Span;
+
+/**
+ * One subword unit produced by a {@link SubwordTokenizer}, carrying both the vocabulary view
+ * (the piece string and its id) and the exact place in the caller's text it came from.
+ *
+ * <p>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}.</p>
+ *
+ * @param piece The piece in the vocabulary's normalized form; never null or empty.
+ * @param id    The vocabulary id of the piece.
+ * @param start The inclusive start offset in the original text.
+ * @param end   The exclusive end offset in the original text; not less than {@code start}.
+ */
+public record SubwordPiece(String piece, int id, int start, int end) {
+
+  /**
+   * Instantiates a {@link SubwordPiece}.
+   *
+   * @throws IllegalArgumentException Thrown if {@code piece} is null or empty, or the span is
+   *     negative or inverted.
+   */
+  public SubwordPiece {
+    if (piece == null || piece.isEmpty()) {
+      throw new IllegalArgumentException("piece must not be null or empty");
+    }
+    if (start < 0 || end < start) {
+      throw new IllegalArgumentException(
+          "The span [" + start + ", " + end + ") must not be negative or inverted.");
+    }
+  }
+
+  /** {@return the original-text span of this piece as a {@link Span}} */
+  public Span span() {
+    return new Span(start, end);
+  }
+}

opennlp-api/src/main/java/opennlp/tools/tokenize/SubwordTokenizer.java

@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package opennlp.tools.tokenize;
+
+import java.util.List;
+
+/**
+ * Splits text into subword units against a fixed vocabulary, reporting for every unit its
+ * vocabulary id and the exact span of the original text it covers.
+ *
+ * <p>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.</p>
+ *
+ * <p>Thread safety is implementation specific.</p>
+ */
+public interface SubwordTokenizer {
+
+  /**
+   * Encodes text into subword pieces.
+   *
+   * @param text The text to encode; must not be null.
+   * @return The pieces in text order; empty when the text contains nothing encodable.
+   * @throws IllegalArgumentException Thrown if {@code text} is null.
+   */
+  List<SubwordPiece> encode(CharSequence text);
+
+  /**
+   * Encodes text into vocabulary ids.
+   *
+   * @param text The text to encode; must not be null.
+   * @return The ids in text order; empty when the text contains nothing encodable.
+   * @throws IllegalArgumentException Thrown if {@code text} is null.
+   */
+  default int[] encodeToIds(CharSequence text) {
+    final List<SubwordPiece> pieces = encode(text);
+    final int[] ids = new int[pieces.size()];
+    for (int i = 0; i < ids.length; i++) {
+      ids[i] = pieces.get(i).id();
+    }
+    return ids;
+  }
+
+  /**
+   * Encodes text into piece strings in the vocabulary's normalized form.
+   *
+   * @param text The text to encode; must not be null.
+   * @return The pieces in text order; empty when the text contains nothing encodable.
+   * @throws IllegalArgumentException Thrown if {@code text} is null.
+   */
+  default String[] encodeToPieces(CharSequence text) {
+    final List<SubwordPiece> pieces = encode(text);
+    final String[] out = new String[pieces.size()];
+    for (int i = 0; i < out.length; i++) {
+      out[i] = pieces.get(i).piece();
+    }
+    return out;
+  }
+}