Refactor for Mistral integration with abstractions

orionpapadakis · orionpapadakis · commit a2ec8cc68ba1 · 2025-06-11T18:38:14.000+03:00
diff --git a/src/main/java/com/example/LlamaApp.java b/src/main/java/com/example/LlamaApp.java
@@ -5,10 +5,8 @@
 import com.example.inference.CategoricalSampler;
 import com.example.inference.Sampler;
 import com.example.inference.ToppSampler;
-import com.example.inference.engine.impl.Model;
-import com.example.inference.engine.impl.Options;
+import com.example.model.Model;
 import com.example.loader.weights.ModelLoader;
-import com.example.loader.weights.State;
 import com.example.tornadovm.FloatArrayUtils;
 import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
 
diff --git a/src/main/java/com/example/Options.java b/src/main/java/com/example/Options.java
@@ -1,4 +1,4 @@
-package com.example.inference.engine.impl;
+package com.example;
 
 import java.io.PrintStream;
 import java.nio.file.Path;
diff --git a/src/main/java/com/example/aot/AOT.java b/src/main/java/com/example/aot/AOT.java
@@ -3,9 +3,9 @@
 import com.example.auxiliary.Timer;
 import com.example.core.model.GGUF;
 import com.example.core.model.tensor.GGMLTensorEntry;
-import com.example.inference.engine.impl.Model;
-import com.example.inference.engine.impl.Options;
-import com.example.inference.engine.impl.llama.Llama;
+import com.example.model.Model;
+import com.example.Options;
+import com.example.model.llama.Llama;
 import com.example.loader.weights.ModelLoader;
 import com.example.loader.weights.Weights;
 
diff --git a/src/main/java/com/example/auxiliary/format/ChatFormat.java b/src/main/java/com/example/auxiliary/format/ChatFormat.java
@@ -4,4 +4,6 @@
 
 public interface ChatFormat {
 
+    int getBeginOfText();
+
 }
diff --git a/src/main/java/com/example/auxiliary/format/LlamaChatFormat.java b/src/main/java/com/example/auxiliary/format/LlamaChatFormat.java
@@ -10,14 +10,14 @@
 
 public class LlamaChatFormat implements ChatFormat {
 
-    final LlamaTokenizer tokenizer;
-    public final int beginOfText;
-    final int endHeader;
-    final int startHeader;
-    final int endOfTurn;
-    final int endOfText;
-    final int endOfMessage;
-    final Set<Integer> stopTokens;
+    protected final LlamaTokenizer tokenizer;
+    protected final int beginOfText;
+    protected final int endHeader;
+    protected final int startHeader;
+    protected final int endOfTurn;
+    protected final int endOfText;
+    protected final int endOfMessage;
+    protected final Set<Integer> stopTokens;
 
     public LlamaChatFormat(LlamaTokenizer tokenizer) {
         this.tokenizer = tokenizer;
@@ -69,6 +69,11 @@ public List<Integer> encodeDialogPrompt(boolean appendAssistantTurn, List<LlamaC
         return tokens;
     }
 
+    @Override
+    public int getBeginOfText() {
+        return beginOfText;
+    }
+
     public record Message(LlamaChatFormat.Role role, String content) {
     }
 
diff --git a/src/main/java/com/example/auxiliary/format/MistralChatFormat.java b/src/main/java/com/example/auxiliary/format/MistralChatFormat.java
@@ -40,6 +40,7 @@ public MistralChatFormat(MistralTokenizer tokenizer) {
         this.middle = specialTokens.getOrDefault("[MIDDLE]", unknownToken);
     }
 
+    @Override
     public int getBeginOfText() { return beginOfText; }
 
     public Set<Integer> getStopTokens() {
diff --git a/src/main/java/com/example/loader/weights/ModelLoader.java b/src/main/java/com/example/loader/weights/ModelLoader.java
@@ -10,12 +10,12 @@
 import com.example.core.model.tensor.Q4_0FloatTensor;
 import com.example.core.model.tensor.Q8_0FloatTensor;
 import com.example.core.types.Pair;
-import com.example.inference.engine.impl.Configuration;
-import com.example.inference.engine.impl.Model;
-import com.example.inference.engine.impl.llama.LlamaConfiguration;
-import com.example.inference.engine.impl.llama.Llama;
-import com.example.inference.engine.impl.mistral.Mistral;
-import com.example.inference.engine.impl.mistral.MistralConfiguration;
+import com.example.model.Configuration;
+import com.example.model.Model;
+import com.example.model.llama.LlamaConfiguration;
+import com.example.model.llama.Llama;
+import com.example.model.mistral.Mistral;
+import com.example.model.mistral.MistralConfiguration;
 import com.example.inference.operation.RoPE;
 import com.example.tokenizer.impl.LlamaTokenizer;
 import com.example.tokenizer.impl.MistralTokenizer;
diff --git a/src/main/java/com/example/loader/weights/State.java b/src/main/java/com/example/loader/weights/State.java
@@ -2,7 +2,7 @@
 
 import com.example.core.model.tensor.ArrayFloatTensor;
 import com.example.core.model.tensor.FloatTensor;
-import com.example.inference.engine.impl.Configuration;
+import com.example.model.Configuration;
 import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
 import uk.ac.manchester.tornado.api.types.arrays.IntArray;
 
diff --git a/src/main/java/com/example/model/Configuration.java b/src/main/java/com/example/model/Configuration.java
@@ -1,4 +1,4 @@
-package com.example.inference.engine.impl;
+package com.example.model;
 
 public interface Configuration {
 
diff --git a/src/main/java/com/example/model/Model.java b/src/main/java/com/example/model/Model.java
@@ -1,6 +1,8 @@
-package com.example.inference.engine.impl;
+package com.example.model;
 
 import com.example.inference.Sampler;
+import com.example.Options;
+import com.example.loader.weights.ModelLoader.ModelType;
 import com.example.loader.weights.State;
 import com.example.loader.weights.Weights;
 import com.example.tokenizer.impl.Tokenizer;
@@ -15,6 +17,8 @@ public interface Model {
     Tokenizer tokenizer();
     Weights weights();
 
+    ModelType getModelType();
+
     List<Integer> generateTokensGPU(State state, int startPosition, List<Integer> promptTokens, Set<Integer> stopTokens,
                                     int maxTokens, Sampler sampler, boolean echo, IntConsumer onTokenGenerated, TornadoVMMasterPlan tornadoVMPlan);
     /**
diff --git a/src/main/java/com/example/model/llama/Llama.java b/src/main/java/com/example/model/llama/Llama.java
@@ -1,12 +1,13 @@
-package com.example.inference.engine.impl.llama;
+package com.example.model.llama;
 
 import com.example.auxiliary.Parallel;
 import com.example.auxiliary.format.LlamaChatFormat;
 import com.example.core.model.tensor.FloatTensor;
 import com.example.inference.Sampler;
-import com.example.inference.engine.impl.Configuration;
-import com.example.inference.engine.impl.Model;
-import com.example.inference.engine.impl.Options;
+import com.example.model.Configuration;
+import com.example.model.Model;
+import com.example.Options;
+import com.example.loader.weights.ModelLoader;
 import com.example.loader.weights.State;
 import com.example.loader.weights.Weights;
 import com.example.tokenizer.impl.LlamaTokenizer;
@@ -207,6 +208,11 @@ private static FloatArray forwardTornadoVM( //
         return tornadoVMMasterPlan.tornadoVMForwardExecuteLayered(position);
     }
 
+    @Override
+    public ModelLoader.ModelType getModelType() {
+        return ModelLoader.ModelType.LLAMA_3;
+    }
+
     @Override
     public List<Integer> generateTokensGPU(State state, int startPosition, List<Integer> promptTokens, Set<Integer> stopTokens,
                                            int maxTokens, Sampler sampler, boolean echo, IntConsumer onTokenGenerated, TornadoVMMasterPlan tornadoVMPlan) {
@@ -299,7 +305,13 @@ public List<Integer> generateTokensGPU(State state, int startPosition, List<Inte
 
     @Override
     public List<Integer> generateTokens(State state,
-                                        int startPosition, List<Integer> promptTokens, Set<Integer> stopTokens, int maxTokens, Sampler sampler, boolean echo, IntConsumer onTokenGenerated) {
+                                        int startPosition,
+                                        List<Integer> promptTokens,
+                                        Set<Integer> stopTokens,
+                                        int maxTokens,
+                                        Sampler sampler,
+                                        boolean echo,
+                                        IntConsumer onTokenGenerated) {
         // Initialize TornadoVM plan if enabled
 
         // Start timing the whole process
@@ -395,7 +407,7 @@ public void runInteractive(Sampler sampler, Options options) {
         State state = null;
         List<Integer> conversationTokens = new ArrayList<>();
         LlamaChatFormat chatFormat = new LlamaChatFormat(getAsLlamaTokenizer());
-        conversationTokens.add(chatFormat.beginOfText);
+        conversationTokens.add(chatFormat.getBeginOfText());
         if (options.systemPrompt() != null) {
             conversationTokens.addAll(chatFormat.encodeMessage(new LlamaChatFormat.Message(LlamaChatFormat.Role.SYSTEM, options.systemPrompt())));
         }
@@ -490,7 +502,7 @@ public void runInstructOnce(Sampler sampler, Options options) {
         TornadoVMMasterPlan tornadoVMPlan =null;
 
         List<Integer> promptTokens = new ArrayList<>();
-        promptTokens.add(chatFormat.beginOfText);
+        promptTokens.add(chatFormat.getBeginOfText());
         if (options.systemPrompt() != null) {
             promptTokens.addAll(chatFormat.encodeMessage(new LlamaChatFormat.Message(LlamaChatFormat.Role.SYSTEM, options.systemPrompt())));
         }
diff --git a/src/main/java/com/example/model/llama/LlamaConfiguration.java b/src/main/java/com/example/model/llama/LlamaConfiguration.java
@@ -1,6 +1,6 @@
-package com.example.inference.engine.impl.llama;
+package com.example.model.llama;
 
-import com.example.inference.engine.impl.Configuration;
+import com.example.model.Configuration;
 
 public record LlamaConfiguration(
         int dim, int hiddenDim, int numberOfLayers, int numberOfHeads,
diff --git a/src/main/java/com/example/model/mistral/Mistral.java b/src/main/java/com/example/model/mistral/Mistral.java
@@ -1,12 +1,13 @@
-package com.example.inference.engine.impl.mistral;
+package com.example.model.mistral;
 
 import com.example.auxiliary.Parallel;
 import com.example.auxiliary.format.MistralChatFormat;
 import com.example.core.model.tensor.FloatTensor;
 import com.example.inference.Sampler;
-import com.example.inference.engine.impl.Configuration;
-import com.example.inference.engine.impl.Model;
-import com.example.inference.engine.impl.Options;
+import com.example.model.Configuration;
+import com.example.model.Model;
+import com.example.Options;
+import com.example.loader.weights.ModelLoader;
 import com.example.loader.weights.State;
 import com.example.loader.weights.Weights;
 import com.example.tokenizer.impl.MistralTokenizer;
@@ -27,6 +28,9 @@
  */
 public record Mistral(MistralConfiguration configuration, Tokenizer tokenizer, Weights weights) implements Model {
 
+    /* For explicit use */
+    private MistralTokenizer getAsMistralTokenizer() { return (MistralTokenizer) tokenizer; }
+
     static void rmsnorm(FloatTensor out, FloatTensor x, FloatBuffer weight, int size, float rmsNormEps) {
         // calculate sum of squares
         float ss = x.reduce(0, size, 0f, (acc, xi) -> acc + xi * xi);
@@ -162,20 +166,20 @@ static FloatTensor forward(Mistral model, State state, int token, int position)
         return state.logits;
     }
 
-    /* For explicit use */
-    private MistralTokenizer getAsMistralTokenizer() {
-        return (MistralTokenizer) tokenizer;
+    @Override
+    public ModelLoader.ModelType getModelType() {
+        return ModelLoader.ModelType.MISTRAL;
     }
 
     @Override
-    public List<Integer> generateTokensGPU(State state, int startPosition, List<Integer> promptTokens, Set<Integer> stopTokens, int maxTokens, Sampler sampler, boolean echo,
-            IntConsumer onTokenGenerated, TornadoVMMasterPlan tornadoVMPlan) {
+    public List<Integer> generateTokensGPU(State state, int startPosition, List<Integer> promptTokens, Set<Integer> stopTokens,
+                                           int maxTokens, Sampler sampler, boolean echo, IntConsumer onTokenGenerated, TornadoVMMasterPlan tornadoVMPlan) {
         throw new UnsupportedOperationException("Mistral.generateTokensGPU is not implemented yet");
     }
 
     @Override
-    public List<Integer> generateTokens(State state, int startPosition, List<Integer> promptTokens, Set<Integer> stopTokens, int maxTokens, Sampler sampler, boolean echo,
-            IntConsumer onTokenGenerated) {
+    public List<Integer> generateTokens(State state, int startPosition, List<Integer> promptTokens, Set<Integer> stopTokens,
+                                        int maxTokens, Sampler sampler, boolean echo, IntConsumer onTokenGenerated) {
         long startNanos = System.nanoTime();
         if (maxTokens < 0 || configuration.contextLength() < maxTokens) {
             maxTokens = configuration.contextLength();
@@ -282,9 +286,8 @@ public void runInteractive(Sampler sampler, Options options) {
 
     @Override
     public void runInstructOnce(Sampler sampler, Options options) {
-        MistralTokenizer mistralTokenizer = getAsMistralTokenizer();
         State state = createNewState();
-        MistralChatFormat chatFormat = new MistralChatFormat(mistralTokenizer);
+        MistralChatFormat chatFormat = new MistralChatFormat(getAsMistralTokenizer());
         List<Integer> promptTokens = new ArrayList<>();
         promptTokens.add(chatFormat.getBeginOfText());
         if (options.suffix() != null) {
@@ -297,9 +300,9 @@ public void runInstructOnce(Sampler sampler, Options options) {
         Set<Integer> stopTokens = chatFormat.getStopTokens();
         IntConsumer tokenConsumer = token -> {
             if (options.stream()) {
-                int tokenType = mistralTokenizer.getTokenType(token);
+                int tokenType = getAsMistralTokenizer().getTokenType(token);
                 if (tokenType == 1 || tokenType == 6) {
-                    System.out.print(mistralTokenizer.decode(List.of(token)));
+                    System.out.print(tokenizer.decode(List.of(token)));
                 }
             }
         };
@@ -317,7 +320,7 @@ public void runInstructOnce(Sampler sampler, Options options) {
             responseTokens.removeLast();
         }
         if (!options.stream()) {
-            String responseText = mistralTokenizer.decode(responseTokens);
+            String responseText = tokenizer.decode(responseTokens);
             System.out.println(responseText);
         }
     }
diff --git a/src/main/java/com/example/model/mistral/MistralConfiguration.java b/src/main/java/com/example/model/mistral/MistralConfiguration.java
@@ -1,6 +1,6 @@
-package com.example.inference.engine.impl.mistral;
+package com.example.model.mistral;
 
-import com.example.inference.engine.impl.Configuration;
+import com.example.model.Configuration;
 
 public record MistralConfiguration(int dim, int hiddenDim, int numberOfLayers, int numberOfHeads, int numberOfKeyValueHeads, int vocabularySize, int contextLength, boolean sharedWeights,
                                    float rmsNormEps, float ropeTheta) implements Configuration {
diff --git a/src/main/java/com/example/tokenizer/impl/LlamaTokenizer.java b/src/main/java/com/example/tokenizer/impl/LlamaTokenizer.java
@@ -11,8 +11,10 @@
 import java.util.stream.IntStream;
 
 public class LlamaTokenizer implements Tokenizer {
+    // general fields
     private final Pattern compiledPattern;
     private final Vocabulary vocabulary;
+    // model-specific fields
     private final Map<Pair<Integer, Integer>, Integer> merges;
     private final Map<String, Integer> specialTokens;
 
@@ -23,10 +25,12 @@ public String regexPattern() {
         return compiledPattern.pattern();
     }
 
+    @Override
     public Map<String, Integer> getSpecialTokens() {
         return specialTokens;
     }
 
+    @Override
     public boolean isSpecialToken(int tokenIndex) {
         return specialTokens.containsValue(tokenIndex);
     }
@@ -220,6 +224,7 @@ public int[] encode(String text) {
         return encodeImpl(sb.toString());
     }
 
+    @Override
     public List<Integer> encodeAsList(String text) {
         StringBuilder sb = new StringBuilder();
         byte[] bytes = text.getBytes(StandardCharsets.UTF_8);
@@ -229,6 +234,7 @@ public List<Integer> encodeAsList(String text) {
         return Arrays.stream(encodeImpl(sb.toString())).boxed().toList();
     }
 
+    @Override
     public String decode(List<Integer> tokens) {
         String decoded = decodeImpl(tokens);
         int[] decodedBytesAsInts = decoded.codePoints().map(BYTE_DECODER::get).toArray();
diff --git a/src/main/java/com/example/tokenizer/impl/MistralTokenizer.java b/src/main/java/com/example/tokenizer/impl/MistralTokenizer.java
@@ -13,8 +13,10 @@
  * <a href="https://github.com/openai/gpt-2/blob/master/src/encoder.py">GPT 2 tokenizer</a>
  */
 public class MistralTokenizer implements Tokenizer {
+    // general fields
     private final Pattern compiledPattern;
     private final Vocabulary vocabulary;
+    // model-specific fields
     private final Map<String, Integer> specialTokens;
     private final int[] tokenType;
     private final int byte0;
diff --git a/src/main/java/com/example/tokenizer/vocabulary/Vocabulary.java b/src/main/java/com/example/tokenizer/vocabulary/Vocabulary.java
@@ -1,5 +1,6 @@
 package com.example.tokenizer.vocabulary;
 
+import java.util.Arrays;
 import java.util.Map;
 import java.util.OptionalInt;
 import java.util.stream.Collectors;
@@ -48,4 +49,17 @@ public int size() {
     public float getScore(int tokenIndex) {
         return scores[tokenIndex];
     }
+
+    @Override
+    public String toString() {
+        StringBuilder sb = new StringBuilder();
+        sb.append("Vocabulary:\n");
+        sb.append("Tokens: ").append(Arrays.toString(tokens)).append("\n");
+        sb.append("Scores: ").append(Arrays.toString(scores)).append("\n");
+        sb.append("Token to Index Map:\n");
+        tokenToIndex.forEach((token, index) ->
+                sb.append("  ").append(token).append(" -> ").append(index).append("\n")
+        );
+        return sb.toString();
+    }
 }
diff --git a/src/main/java/com/example/tornadovm/TornadoVMLayerPlanner.java b/src/main/java/com/example/tornadovm/TornadoVMLayerPlanner.java
@@ -1,10 +1,8 @@
 package com.example.tornadovm;
 
 import com.example.auxiliary.Tuple2;
-import com.example.inference.engine.impl.Configuration;
-import com.example.inference.engine.impl.Model;
-import com.example.inference.engine.impl.llama.LlamaConfiguration;
-import com.example.inference.engine.impl.llama.Llama;
+import com.example.model.Configuration;
+import com.example.model.Model;
 import com.example.loader.weights.State;
 import com.example.loader.weights.Weights;
 import uk.ac.manchester.tornado.api.GridScheduler;
diff --git a/src/main/java/com/example/tornadovm/TornadoVMMasterPlan.java b/src/main/java/com/example/tornadovm/TornadoVMMasterPlan.java

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-package com.example.inference.engine.impl;`
	`1`	`+package com.example;`
`2`	`2`
`3`	`3`	`import java.io.PrintStream;`
`4`	`4`	`import java.nio.file.Path;`
Original file line number	Diff line number	Diff line change
`@@ -4,4 +4,6 @@`
`4`	`4`
`5`	`5`	`public interface ChatFormat {`
`6`	`6`
	`7`	`+ int getBeginOfText();`
	`8`	`+`
`7`	`9`	`}`
Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,7 @@ public MistralChatFormat(MistralTokenizer tokenizer) {`
`40`	`40`	`this.middle = specialTokens.getOrDefault("[MIDDLE]", unknownToken);`
`41`	`41`	`}`
`42`	`42`
	`43`	`+ @Override`
`43`	`44`	`public int getBeginOfText() { return beginOfText; }`
`44`	`45`
`45`	`46`	`public Set<Integer> getStopTokens() {`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-package com.example.inference.engine.impl;`
	`1`	`+package com.example.model;`
`2`	`2`
`3`	`3`	`public interface Configuration {`
`4`	`4`