Refactor to separate classes

8 files changed, 279 insertions, 203 deletions

diff --git a/linguistics/src/main/java/com/yahoo/language/sentencepiece/Model.java b/linguistics/src/main/java/com/yahoo/language/sentencepiece/Model.java
new file mode 100644
index 00000000000..1e2d5c87f66
--- /dev/null
+++ b/linguistics/src/main/java/com/yahoo/language/sentencepiece/Model.java
@@ -0,0 +1,62 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+//
+
+package com.yahoo.language.sentencepiece;
+
+import com.yahoo.io.IOUtils;
+import com.yahoo.language.Language;
+import sentencepiece.SentencepieceModel;
+
+import java.io.IOException;
+import java.nio.file.Path;
+
+/**
+ * A SentencePiece model
+ *
+ * @author bratseth
+ */
+final class Model {
+
+    final Path source;
+    final Language language;
+    final float minScore;
+    final float maxScore;
+    final Trie tokens = new Trie();
+
+    Model(Language language, Path path) {
+        try {
+            this.source = path;
+            this.language = language;
+            var sp = SentencepieceModel.ModelProto.parseFrom(IOUtils.readFileBytes(path.toFile()));
+            float minScore = Float.MAX_VALUE;
+            float maxScore = Float.MIN_VALUE;
+            for (int i = 0; i < sp.getPiecesCount(); i++) {
+                var piece = sp.getPieces(i);
+                tokens.add(toTokenType(piece.getType()), i, piece.getPiece(), piece.getScore());
+                minScore = Math.min(piece.getScore(), minScore);
+                maxScore = Math.max(piece.getScore(), maxScore);
+            }
+            this.minScore = minScore;
+            this.maxScore = maxScore;
+        } catch (IOException e) {
+            throw new IllegalArgumentException("Could not read a SentencePiece model from " + path, e);
+        }
+    }
+
+    private static TokenType toTokenType(SentencepieceModel.ModelProto.SentencePiece.Type type) {
+        switch (type) {
+            case USER_DEFINED : return TokenType.userDefined;
+            case UNKNOWN : return TokenType.unknown;
+            case NORMAL : return TokenType.text;
+            case CONTROL : return TokenType.control;
+            case UNUSED : return TokenType.unused;
+            default : throw new IllegalArgumentException("Unknkown token type " + type);
+        }
+    }
+
+    @Override
+    public String toString() {
+        return "SentencePiece model for " + language + ": '" + source + "'";
+    }
+
+}
diff --git a/linguistics/src/main/java/com/yahoo/language/sentencepiece/ResultBuilder.java b/linguistics/src/main/java/com/yahoo/language/sentencepiece/ResultBuilder.java
new file mode 100644
index 00000000000..2141505374c
--- /dev/null
+++ b/linguistics/src/main/java/com/yahoo/language/sentencepiece/ResultBuilder.java
@@ -0,0 +1,47 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.language.sentencepiece;
+
+/**
+ * Builds a result from a sentencepiece tokenization by being called for each segment in reverse
+ *
+ * @param <RESULTTYPE> the type of result this produces
+ * @author bratseth
+ */
+abstract class ResultBuilder<RESULTTYPE> {
+
+    private final RESULTTYPE result;
+
+    ResultBuilder(RESULTTYPE result) {
+        this.result = result;
+    }
+
+    /** Called for each segment, starting from the last and working backwards */
+    abstract void add(int start, int end, SentencePieceAlgorithm.SegmentEnd[] segmentEnds);
+
+    RESULTTYPE result() {return result;}
+
+    void build(String input, SentencePieceAlgorithm.SegmentEnd[] segmentEnds, boolean collapseUnknowns) {
+        if (collapseUnknowns) {
+            int segmentEnd = input.length();
+            int collapsedSegmentEnd = segmentEnd;
+            while (segmentEnd > 0) {
+                if (segmentEnds[segmentEnd].type != TokenType.unknown ) {
+                    if (collapsedSegmentEnd != segmentEnd) { // We have deferred an unknown collapsed segment
+                        add(segmentEnd, collapsedSegmentEnd, segmentEnds);
+                    }
+                    add(segmentEnds[segmentEnd].segmentStart, segmentEnd, segmentEnds);
+                    collapsedSegmentEnd = segmentEnds[segmentEnd].segmentStart;
+                }
+                segmentEnd = segmentEnds[segmentEnd].segmentStart;
+            }
+        }
+        else {
+            int segmentEnd = input.length();
+            while (segmentEnd > 0) {
+                add(segmentEnds[segmentEnd].segmentStart, segmentEnd, segmentEnds);
+                segmentEnd = segmentEnds[segmentEnd].segmentStart;
+            }
+        }
+    }
+
+}
diff --git a/linguistics/src/main/java/com/yahoo/language/sentencepiece/Scoring.java b/linguistics/src/main/java/com/yahoo/language/sentencepiece/Scoring.java
new file mode 100644
index 00000000000..6c8560abee7
--- /dev/null
+++ b/linguistics/src/main/java/com/yahoo/language/sentencepiece/Scoring.java
@@ -0,0 +1,17 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.language.sentencepiece;
+
+/**
+ * The scoring strategy to use for picking segments
+ *
+ * @author bratseth
+ */
+public enum Scoring {
+
+    /** Find the segmentation that has the highest score */
+    highestScore,
+
+    /** Find the segmentation that has the fewest segments, resolve ties by score sum */
+    fewestSegments
+
+}
diff --git a/linguistics/src/main/java/com/yahoo/language/sentencepiece/SentencePieceAlgorithm.java b/linguistics/src/main/java/com/yahoo/language/sentencepiece/SentencePieceAlgorithm.java
new file mode 100644
index 00000000000..0446b13f0ee
--- /dev/null
+++ b/linguistics/src/main/java/com/yahoo/language/sentencepiece/SentencePieceAlgorithm.java
@@ -0,0 +1,92 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+//
+
+package com.yahoo.language.sentencepiece;
+
+/**
+ * SentencePiece algorithm implementation
+ *
+ * @author bratseth
+ */
+class SentencePieceAlgorithm {
+
+    // TODO: Support characters beyond BMP
+
+    static final char spaceSymbol = '▁';
+
+    private final boolean collapseUnknowns;
+    private final Scoring scoring;
+
+    SentencePieceAlgorithm(boolean collapseUnknowns, Scoring scoring) {
+        this.collapseUnknowns = collapseUnknowns;
+        this.scoring = scoring;
+    }
+
+    public <RESULTTYPE> void segment(String input, ResultBuilder<RESULTTYPE> resultBuilder, Model model) {
+        SegmentEnd[] segmentEnds = new SegmentEnd[input.length() + 1];
+        segmentEnds[0] = new SegmentEnd(TokenType.unknown, 0, 0, 0, 0);
+        int start = 0;
+        while (start < input.length()) { // segment from this position to the end of the text
+            Trie.Node node = model.tokens.root;
+            int characterPosition = start;
+            while (node != null && characterPosition < input.length()) { // traverse the trie one character at the time from this position
+                node = node.children.get(input.charAt(characterPosition++));
+                int length = characterPosition - start;
+                if (node != null && node.isToken() && node.type != TokenType.unused) {
+                    float score = node.type == TokenType.userDefined ? (length * model.maxScore - 0.1f) : node.score;
+                    addSegment(TokenType.text, node.id, start, characterPosition, score, segmentEnds);
+                }
+                else if (length == 1) { // add an 'unknown' length 1 token to make the next position reachable
+                    addSegment(TokenType.unknown, 0, start, start + 1, model.minScore - 10.0f, segmentEnds);
+                }
+            }
+            start++;
+        }
+        resultBuilder.build(input, segmentEnds, collapseUnknowns);
+    }
+
+    private void addSegment(TokenType type, int id, int start, int end, float score, SegmentEnd[] segmentEnds) {
+        if (segmentEnds[end] == null ||
+            segmentEnds[start].scoreWith(score) > segmentEnds[end].score()) {
+            segmentEnds[end] = new SegmentEnd(type, id,
+                                              segmentEnds[start].pathScoreSum + score,
+                                              segmentEnds[start].pathSegmentCount + 1,
+                                              start);
+        }
+    }
+
+    final class SegmentEnd {
+
+        final TokenType type;
+        final int id;
+        final float pathScoreSum;
+        final int pathSegmentCount;
+        final int segmentStart;
+
+        SegmentEnd(TokenType type, int id, float pathScoreSum, int pathSegmentCount, int segmentStart) {
+            this.type = type;
+            this.id = id;
+            this.pathScoreSum = pathScoreSum;
+            this.pathSegmentCount = pathSegmentCount;
+            this.segmentStart = segmentStart;
+        }
+
+        public float score() {
+            switch (scoring) {
+                case fewestSegments: return 1f / pathSegmentCount * 10_000_000 + pathScoreSum;
+                case highestScore: return pathScoreSum;
+                default : throw new IllegalArgumentException("Unknown scoring " + scoring);
+            }
+        }
+
+        public float scoreWith(float additionalSegmentScore) {
+            switch (scoring) {
+                case fewestSegments: return 1f / (pathSegmentCount + 1) * 10_000_000 + (pathScoreSum + additionalSegmentScore );
+                case highestScore: return pathScoreSum + additionalSegmentScore;
+                default : throw new IllegalArgumentException("Unknown scoring " + scoring);
+            }
+        }
+
+    }
+
+}
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 74ed79b267b..53e8513a611 100644
--- a/linguistics/src/main/java/com/yahoo/language/sentencepiece/SentencePieceEncoder.java
+++ b/linguistics/src/main/java/com/yahoo/language/sentencepiece/SentencePieceEncoder.java
@@ -4,16 +4,13 @@ package com.yahoo.language.sentencepiece;
 
 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;
 import com.yahoo.tensor.TensorType;
-import sentencepiece.SentencepieceModel;
 
-import java.io.IOException;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -33,32 +30,17 @@ import java.util.stream.Collectors;
 @Beta
 public class SentencePieceEncoder implements Segmenter, Encoder {
 
-    // TODO: Support characters beyond BMP
-    enum TokenType { text, control, userDefined, unknown, unused }
-
-    /** The scoring strategy to use for picking segments */
-    public enum Scoring {
-        /** Find the segmentation that has the highest score */
-        highestScore,
-        /** Find the segmentation that has the fewest segments, resolve ties by score sum */
-        fewestSegments
-    }
-
-    private static final char spaceSymbol = '▁';
-
-    private final boolean collapseUnknowns;
-    private final Scoring scoring;
-
     private final Map<Language, Model> models;
 
+    private final SentencePieceAlgorithm algorithm;
+
     @Inject
     public SentencePieceEncoder(SentencePieceConfig config) {
         this(new Builder(config));
     }
 
     public SentencePieceEncoder(Builder builder) {
-        collapseUnknowns = builder.getCollapseUnknowns();
-        scoring = builder.getScoring();
+        algorithm = new SentencePieceAlgorithm(builder.collapseUnknowns, builder.getScoring());
 
         models = builder.getModels().entrySet()
                         .stream()
@@ -79,7 +61,7 @@ public class SentencePieceEncoder implements Segmenter, Encoder {
     public List<String> segment(String rawInput, Language language) {
         String input = normalize(rawInput);
         var resultBuilder = new ResultBuilder<List<String>>(new ArrayList<>()) {
-            public void add(int segmentStart, int segmentEnd, SegmentEnd[] segmentEnds) {
+            public void add(int segmentStart, int segmentEnd, SentencePieceAlgorithm.SegmentEnd[] segmentEnds) {
                 result().add(input.substring(segmentStart, segmentEnd));
             }
         };
@@ -98,7 +80,7 @@ public class SentencePieceEncoder implements Segmenter, Encoder {
     @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) {
+            public void add(int segmentStart, int segmentEnd, SentencePieceAlgorithm.SegmentEnd[] segmentEnds) {
                 result().add(segmentEnds[segmentEnd].id);
             }
         };
@@ -148,29 +130,10 @@ public class SentencePieceEncoder implements Segmenter, Encoder {
         }
     }
 
-    private <RESULTTYPE> void segment(String input, Language language, ResultBuilder<RESULTTYPE> resultBuilder) {
+    private <RESULTTYPE> void segment(String input, Language language,
+                                      ResultBuilder<RESULTTYPE> resultBuilder) {
         Model model = resolveFrom(language);
-
-        SegmentEnd[] segmentEnds = new SegmentEnd[input.length() + 1];
-        segmentEnds[0] = new SegmentEnd(TokenType.unknown, 0, 0, 0, 0);
-        int start = 0;
-        while (start < input.length()) { // segment from this position to the end of the text
-            Trie.Node node = model.tokens.root;
-            int characterPosition = start;
-            while (node != null && characterPosition < input.length()) { // traverse the trie one character at the time from this position
-                node = node.children.get(input.charAt(characterPosition++));
-                int length = characterPosition - start;
-                if (node != null && node.isToken() && node.type != TokenType.unused) {
-                    float score = node.type == TokenType.userDefined ? (length * model.maxScore - 0.1f) : node.score;
-                    addSegment(TokenType.text, node.id, start, characterPosition, score, segmentEnds);
-                }
-                else if (length == 1) { // add an 'unknown' length 1 token to make the next position reachable
-                    addSegment(TokenType.unknown, 0, start, start + 1, model.minScore - 10.0f, segmentEnds);
-                }
-            }
-            start++;
-        }
-        createResult(input, segmentEnds, resultBuilder);
+        algorithm.segment(input, resultBuilder, model);
     }
 
     private Model resolveFrom(Language language) {
@@ -180,88 +143,6 @@ public class SentencePieceEncoder implements Segmenter, Encoder {
         throw new IllegalArgumentException("No SentencePiece model for language " + language + " is configured");
     }
 
-    private void addSegment(TokenType type, int id, int start, int end, float score, SegmentEnd[] segmentEnds) {
-        if (segmentEnds[end] == null ||
-            segmentEnds[start].scoreWith(score) > segmentEnds[end].score()) {
-            segmentEnds[end] = new SegmentEnd(type, id,
-                                                      segmentEnds[start].pathScoreSum + score,
-                                                      segmentEnds[start].pathSegmentCount + 1,
-                                                      start);
-        }
-    }
-
-    private <RESULTTYPE> void createResult(String input, SegmentEnd[] segmentEnds, ResultBuilder<RESULTTYPE> resultBuilder) {
-        if (collapseUnknowns) {
-            int segmentEnd = input.length();
-            int collapsedSegmentEnd = segmentEnd;
-            while (segmentEnd > 0) {
-                if (segmentEnds[segmentEnd].type != TokenType.unknown ) {
-                    if (collapsedSegmentEnd != segmentEnd) { // We have deferred an unknown collapsed segment
-                        resultBuilder.add(segmentEnd, collapsedSegmentEnd, segmentEnds);
-                    }
-                    resultBuilder.add(segmentEnds[segmentEnd].segmentStart, segmentEnd, segmentEnds);
-                    collapsedSegmentEnd = segmentEnds[segmentEnd].segmentStart;
-                }
-                segmentEnd = segmentEnds[segmentEnd].segmentStart;
-            }
-        }
-        else {
-            int segmentEnd = input.length();
-            while (segmentEnd > 0) {
-                resultBuilder.add(segmentEnds[segmentEnd].segmentStart, segmentEnd, segmentEnds);
-                segmentEnd = segmentEnds[segmentEnd].segmentStart;
-            }
-        }
-    }
-
-    private static abstract class ResultBuilder<RESULTTYPE> {
-
-        private final RESULTTYPE result;
-
-        ResultBuilder(RESULTTYPE result) {
-            this.result = result;
-        }
-
-        abstract void add(int start, int end, SegmentEnd[] segmentEnds);
-
-        RESULTTYPE result() { return result; }
-
-    }
-
-    private final class SegmentEnd {
-
-        final TokenType type;
-        final int id;
-        final float pathScoreSum;
-        final int pathSegmentCount;
-        final int segmentStart;
-
-        SegmentEnd(TokenType type, int id, float pathScoreSum, int pathSegmentCount, int segmentStart) {
-            this.type = type;
-            this.id = id;
-            this.pathScoreSum = pathScoreSum;
-            this.pathSegmentCount = pathSegmentCount;
-            this.segmentStart = segmentStart;
-        }
-
-        public float score() {
-            switch (scoring) {
-                case fewestSegments: return 1f / pathSegmentCount * 10_000_000 + pathScoreSum;
-                case highestScore: return pathScoreSum;
-                default : throw new IllegalArgumentException("Unknown scoring " + scoring);
-            }
-        }
-
-        public float scoreWith(float additionalSegmentScore) {
-            switch (scoring) {
-                case fewestSegments: return 1f / (pathSegmentCount + 1) * 10_000_000 + (pathScoreSum + additionalSegmentScore );
-                case highestScore: return pathScoreSum + additionalSegmentScore;
-                default : throw new IllegalArgumentException("Unknown scoring " + scoring);
-            }
-        }
-
-    }
-
     public String normalize(String s) {
         StringBuilder b = new StringBuilder(s.length() + 1);
         boolean queuedSpace = true; // Always start by one space
@@ -272,7 +153,7 @@ public class SentencePieceEncoder implements Segmenter, Encoder {
             }
             else {
                 if (queuedSpace) {
-                    b.append(spaceSymbol);
+                    b.append(SentencePieceAlgorithm.spaceSymbol);
                     queuedSpace = false;
                 }
                 b.append(c);
@@ -281,79 +162,6 @@ public class SentencePieceEncoder implements Segmenter, Encoder {
         return b.toString();
     }
 
-    private static TokenType toTokenType(SentencepieceModel.ModelProto.SentencePiece.Type type) {
-        switch (type) {
-            case USER_DEFINED : return TokenType.userDefined;
-            case UNKNOWN : return TokenType.unknown;
-            case NORMAL : return TokenType.text;
-            case CONTROL : return TokenType.control;
-            case UNUSED : return TokenType.unused;
-            default : throw new IllegalArgumentException("Unknkown token type " + type);
-        }
-    }
-
-    private static class Trie {
-
-        final Node root = new Node();
-
-        void add(TokenType type, int id, String word, float score) {
-            Node current = root;
-            for (char l : word.toCharArray())
-                current = current.children.computeIfAbsent(l, c -> new Node());
-            current.type = type;
-            current.id = id;
-            current.score = score;
-        }
-
-        static class Node {
-
-            Integer id;
-            TokenType type;
-            Float score;
-            private final Map<Character, Node> children = new HashMap<>();
-
-            boolean isToken() { return type != null; }
-
-        }
-
-    }
-
-    private static final class Model {
-
-        final Path source;
-        final Language language;
-        final float minScore;
-        final float maxScore;
-        final Trie tokens = new Trie();
-
-        Model(Language language, Path path) {
-            try {
-                this.source = path;
-                this.language = language;
-                var sp = SentencepieceModel.ModelProto.parseFrom(IOUtils.readFileBytes(path.toFile()));
-                float minScore = Float.MAX_VALUE;
-                float maxScore = Float.MIN_VALUE;
-                for (int i = 0; i < sp.getPiecesCount(); i++) {
-                    var piece = sp.getPieces(i);
-                    tokens.add(toTokenType(piece.getType()), i, piece.getPiece(), piece.getScore());
-                    minScore = Math.min(piece.getScore(), minScore);
-                    maxScore = Math.max(piece.getScore(), maxScore);
-                }
-                this.minScore = minScore;
-                this.maxScore = maxScore;
-            }
-            catch (IOException e) {
-                throw new IllegalArgumentException("Could not read a SentencePiece model from " + path, e);
-            }
-        }
-
-        @Override
-        public String toString() {
-            return "SentencePiece model for " + language + ": '" + source + "'";
-        }
-
-    }
-
     public static class Builder {
 
         private final Map<Language, Path> models = new HashMap<>();
diff --git a/linguistics/src/main/java/com/yahoo/language/sentencepiece/TokenType.java b/linguistics/src/main/java/com/yahoo/language/sentencepiece/TokenType.java
new file mode 100644
index 00000000000..ac78b5c78a6
--- /dev/null
+++ b/linguistics/src/main/java/com/yahoo/language/sentencepiece/TokenType.java
@@ -0,0 +1,13 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+//
+
+package com.yahoo.language.sentencepiece;
+
+/**
+ * SentencePiece token types
+ *
+ * @author bratseth
+ */
+enum TokenType {
+    text, control, userDefined, unknown, unused
+}
diff --git a/linguistics/src/main/java/com/yahoo/language/sentencepiece/Trie.java b/linguistics/src/main/java/com/yahoo/language/sentencepiece/Trie.java
new file mode 100644
index 00000000000..ca54d94f88a
--- /dev/null
+++ b/linguistics/src/main/java/com/yahoo/language/sentencepiece/Trie.java
@@ -0,0 +1,38 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+//
+
+package com.yahoo.language.sentencepiece;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * A simple trie for sentencepiece token lookup
+ *
+ * @author bratseth
+ */
+public class Trie {
+
+    final Node root = new Node();
+
+    void add(TokenType type, int id, String word, float score) {
+        Node current = root;
+        for (char l : word.toCharArray())
+            current = current.children.computeIfAbsent(l, c -> new Node());
+        current.type = type;
+        current.id = id;
+        current.score = score;
+    }
+
+    static class Node {
+
+        Integer id;
+        TokenType type;
+        Float score;
+        final Map<Character, Node> children = new HashMap<>();
+
+        boolean isToken() { return type != null; }
+
+    }
+
+}
diff --git a/linguistics/src/test/java/com/yahoo/language/sentencepiece/SentencePieceTest.java b/linguistics/src/test/java/com/yahoo/language/sentencepiece/SentencePieceTest.java
index 5b77324a6fc..d60d7386d4b 100644
--- a/linguistics/src/test/java/com/yahoo/language/sentencepiece/SentencePieceTest.java
+++ b/linguistics/src/test/java/com/yahoo/language/sentencepiece/SentencePieceTest.java
@@ -3,7 +3,6 @@
 package com.yahoo.language.sentencepiece;
 
 import com.yahoo.language.Language;
-import com.yahoo.tensor.Tensor;
 import org.junit.Test;
 
 import java.io.File;
@@ -69,7 +68,7 @@ public class SentencePieceTest {
     public void testHighestScore() {
         var tester = new SentencePieceTester(new SentencePieceEncoder.Builder()
                                                      .addDefaultModel(new File("src/test/models/sentencepiece/en.wiki.bpe.vs10000.model").toPath())
-                                                     .setScoring(SentencePieceEncoder.Scoring.highestScore));
+                                                     .setScoring(Scoring.highestScore));
         tester.assertSegmented("h", "▁h");
         tester.assertSegmented("he", "▁he");
         tester.assertSegmented("hel", "▁h", "el");