Refactor Model design. Abandon Records, adopt interface with abstract base class hierarchy

Author: orionpapadakis

src/main/java/com/example/model/AbstractModel.java

@@ -0,0 +1,47 @@
+package com.example.model;
+
+import com.example.inference.weights.Weights;
+import com.example.model.format.ChatFormat;
+import com.example.tokenizer.impl.Tokenizer;
+import com.example.tornadovm.TornadoVMMasterPlan;
+
+public abstract class AbstractModel implements Model {
+
+    protected Tokenizer tokenizer;
+    protected Weights weights;
+    protected ChatFormat chatFormat;
+    /**
+     * Represents the master plan for the TornadoVM execution in the context of the model.
+     * This variable is used to manage the execution flow or strategy within the TornadoVM environment.
+     * <p>
+     * Initialized *only* when Tornado is used, via {@link TornadoVMMasterPlan#initializeTornadoVMPlan}.
+     * </p>
+     */
+    protected TornadoVMMasterPlan plan;
+
+    protected AbstractModel(Tokenizer tokenizer, Weights weights, ChatFormat chatFormat, TornadoVMMasterPlan plan) {
+        this.tokenizer = tokenizer;
+        this.weights = weights;
+        this.chatFormat = chatFormat;
+        this.plan = plan;
+    }
+
+    // Common methods across models
+
+    public Weights weights() {
+        return weights;
+    }
+
+    public ChatFormat chatFormat() {
+        return chatFormat;
+    }
+
+    public TornadoVMMasterPlan tornadoVMPlan() {
+        return plan;
+    }
+
+    public void setTornadoVMPlan(TornadoVMMasterPlan plan) {
+        this.plan = plan;
+    }
+
+}

src/main/java/com/example/model/Model.java

@@ -29,6 +29,10 @@ public interface Model {
 
     ChatFormat chatFormat();
 
+    TornadoVMMasterPlan tornadoVMPlan();
+
+    void setTornadoVMPlan(TornadoVMMasterPlan plan);
+
     ModelType getModelType();
 
     State createNewState();

src/main/java/com/example/model/llama/Llama.java

@@ -6,7 +6,7 @@
 import com.example.inference.state.LlamaState;
 import com.example.inference.state.State;
 import com.example.inference.weights.Weights;
-import com.example.model.Model;
+import com.example.model.AbstractModel;
 import com.example.model.ModelType;
 import com.example.model.format.ChatFormat;
 import com.example.tokenizer.impl.LlamaTokenizer;
@@ -16,10 +16,22 @@
 import java.util.Set;
 import java.util.function.IntConsumer;
 
-public record Llama(LlamaConfiguration configuration, Tokenizer tokenizer, Weights weights, ChatFormat chatFormat) implements Model {
+public class Llama extends AbstractModel {
 
-    /* For explicit use */
-    private LlamaTokenizer getAsLlamaTokenizer() {
+    LlamaConfiguration configuration;
+
+    public Llama(LlamaConfiguration configuration, Tokenizer tokenizer, Weights weights, ChatFormat chatFormat) {
+        super(tokenizer, weights, chatFormat, null);
+        this.configuration = configuration;
+    }
+
+    @Override
+    public LlamaConfiguration configuration() {
+        return configuration;
+    }
+
+    @Override
+    public LlamaTokenizer tokenizer() {
         return (LlamaTokenizer) tokenizer;
     }
 

src/main/java/com/example/model/mistral/Mistral.java

@@ -6,7 +6,7 @@
 import com.example.inference.state.LlamaState;
 import com.example.inference.state.State;
 import com.example.inference.weights.Weights;
-import com.example.model.Model;
+import com.example.model.AbstractModel;
 import com.example.model.ModelType;
 import com.example.model.format.ChatFormat;
 import com.example.tokenizer.impl.MistralTokenizer;
@@ -16,11 +16,23 @@
 import java.util.Set;
 import java.util.function.IntConsumer;
 
-public record Mistral(MistralConfiguration configuration, Tokenizer tokenizer, Weights weights, ChatFormat chatFormat) implements Model {
+public class Mistral extends AbstractModel {
 
-    /* For explicit use */
-    private MistralTokenizer getAsMistralTokenizer() {
-        return (MistralTokenizer) tokenizer;
+    MistralConfiguration configuration;
+
+    public Mistral(MistralConfiguration configuration, Tokenizer tokenizer, Weights weights, ChatFormat chatFormat) {
+        super(tokenizer, weights, chatFormat, null);
+        this.configuration = configuration;
+    }
+
+    @Override
+    public MistralConfiguration configuration() {
+        return configuration;
+    }
+
+    @Override
+    public MistralTokenizer tokenizer() {
+        return (MistralTokenizer)tokenizer;
     }
 
     @Override

src/main/java/com/example/model/qwen3/Qwen3.java

@@ -6,22 +6,38 @@
 import com.example.inference.state.Qwen3State;
 import com.example.inference.state.State;
 import com.example.inference.weights.Weights;
-import com.example.model.Model;
+import com.example.model.AbstractModel;
 import com.example.model.ModelType;
 import com.example.model.format.ChatFormat;
+import com.example.tokenizer.impl.Qwen3Tokenizer;
 import com.example.tokenizer.impl.Tokenizer;
 
 import java.util.List;
 import java.util.Set;
 import java.util.function.IntConsumer;
 
-public record Qwen3(Qwen3Configuration configuration, Tokenizer tokenizer, Weights weights, ChatFormat chatFormat) implements Model {
+public class Qwen3 extends AbstractModel {
+
+    Qwen3Configuration configuration;
+
+    public Qwen3(Qwen3Configuration configuration, Tokenizer tokenizer, Weights weights, ChatFormat chatFormat) {
+        super(tokenizer, weights, chatFormat, null);
+        this.configuration = configuration;
+    }
+
+    public Qwen3Configuration configuration() {
+        return configuration;
+    }
 
     @Override
     public ModelType getModelType() {
         return ModelType.QWEN_3;
     }
 
+    public Qwen3Tokenizer tokenizer() {
+        return (Qwen3Tokenizer) tokenizer;
+    }
+
     @Override
     public State createNewState() {
         State state = new Qwen3State(configuration(), -1);