Encoder interface

3 files changed, 55 insertions, 4 deletions

diff --git a/linguistics/src/main/java/com/yahoo/language/process/Encoder.java b/linguistics/src/main/java/com/yahoo/language/process/Encoder.java
new file mode 100644
index 00000000000..91de16f669b
--- /dev/null
+++ b/linguistics/src/main/java/com/yahoo/language/process/Encoder.java
@@ -0,0 +1,39 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.language.process;
+
+import com.yahoo.language.Language;
+import com.yahoo.tensor.Tensor;
+import com.yahoo.tensor.TensorType;
+
+import java.util.List;
+
+/**
+ * An encoder converts a text string to a tensor or list of tokens
+ *
+ * @author bratseth
+ */
+public interface Encoder {
+
+    /**
+     * Encodes text into tokens in a list of ids.
+     *
+     * @param text the text to encode
+     * @param language the language of the text, or UNKNOWN to use language independent encoding
+     * @return the text encoded to a list of segment ids
+     * @throws IllegalArgumentException if the language is not supported by this encoder
+     */
+    List<Integer> encode(String text, Language language);
+
+    /**
+     * Encodes text into tokens in a tensor.
+     * The information contained in the encoding may depend on the tensor type.
+     *
+     * @param text the text to encode
+     * @param language the language of the text, or UNKNOWN to use language independent encoding
+     * @param tensorType the type of the ttensor to be returned
+     * @return the tex encoded into a tensor of the supplied type
+     * @throws IllegalArgumentException if the language or tensor type is not supported by this encoder
+     */
+    Tensor encode(String text, Language language, TensorType tensorType);
+
+}
diff --git a/linguistics/src/main/java/com/yahoo/language/process/Stemmer.java b/linguistics/src/main/java/com/yahoo/language/process/Stemmer.java
index da8a73407ff..a2d0d0a84c9 100644
--- a/linguistics/src/main/java/com/yahoo/language/process/Stemmer.java
+++ b/linguistics/src/main/java/com/yahoo/language/process/Stemmer.java
@@ -6,9 +6,9 @@ import com.yahoo.language.Language;
 import java.util.List;
 
 /**
- * <p>Interface providing stemming of single words.</p>
+ * Interface providing stemming of single words.
  *
- * @author <a href="mailto:mathiasm@yahoo-inc.com">Mathias Mølster Lidal</a>
+ * @author Mathias Mølster Lidal
  */
 public interface Stemmer {
 
diff --git a/linguistics/src/main/java/com/yahoo/language/sentencepiece/SentencePieceEncoder.java b/linguistics/src/main/java/com/yahoo/language/sentencepiece/SentencePieceEncoder.java
index c7b131cc439..74ed79b267b 100644
--- a/linguistics/src/main/java/com/yahoo/language/sentencepiece/SentencePieceEncoder.java
+++ b/linguistics/src/main/java/com/yahoo/language/sentencepiece/SentencePieceEncoder.java
@@ -6,6 +6,7 @@ import com.google.common.annotations.Beta;
 import com.google.inject.Inject;
 import com.yahoo.io.IOUtils;
 import com.yahoo.language.Language;
+import com.yahoo.language.process.Encoder;
 import com.yahoo.language.process.Segmenter;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorAddress;
@@ -30,7 +31,7 @@ import java.util.stream.Collectors;
  * @author bratseth
  */
 @Beta
-public class SentencePieceEncoder implements Segmenter {
+public class SentencePieceEncoder implements Segmenter, Encoder {
 
     // TODO: Support characters beyond BMP
     enum TokenType { text, control, userDefined, unknown, unused }
@@ -94,6 +95,7 @@ public class SentencePieceEncoder implements Segmenter {
      * @param language the model to use, or Language.UNKNOWN to use the default model if any
      * @return the list of zero or more token ids resulting from segmenting the input text
      */
+    @Override
     public List<Integer> encode(String rawInput, Language language) {
         var resultBuilder = new ResultBuilder<List<Integer>>(new ArrayList<>()) {
             public void add(int segmentStart, int segmentEnd, SegmentEnd[] segmentEnds) {
@@ -106,8 +108,18 @@ public class SentencePieceEncoder implements Segmenter {
     }
 
     /**
-     * Encodes directly to a tensor.
+     * <p>Encodes directly to a tensor.</p>
+     *
+     * <p>If the tensor type is indexed 1-d (bound or unbound) this will return a tensor containing the token ids in the order
+     * they were encountered in the text. If the dimension is bound and too large it will be zero padded, if too small
+     * it will be truncated.</p>
+     *
+     * <p>If the tensor type is1-d sparse this will return a tensor containing the token strings as keys and the token
+     * position as value.</p>
+     *
+     * <p>If the tensor is any other type IllegalArgumentException is thrown.</p>
      */
+    @Override
     public Tensor encode(String rawInput, Language language, TensorType type) {
         if (type.dimensions().size() == 1 && type.dimensions().get(0).isIndexed()) {
             // Build to a list first since we can't reverse a tensor builder