Refactor ModelLoader: move loadModel() from LlamaApp to ModelLoader class

Author: orionpapadakis

src/main/java/org/beehive/gpullama3/LlamaApp.java

@@ -16,42 +16,13 @@
 import java.util.random.RandomGeneratorFactory;
 
 import static org.beehive.gpullama3.inference.sampler.Sampler.createSampler;
+import static org.beehive.gpullama3.model.loader.ModelLoader.loadModel;
+
 public class LlamaApp {
     // Configuration flags for hardware acceleration and optimizations
     public static final boolean USE_VECTOR_API = Boolean.parseBoolean(System.getProperty("llama.VectorAPI", "true"));   // Enable Java Vector API for CPU acceleration
-    public static final boolean USE_AOT = Boolean.parseBoolean(System.getProperty("llama.AOT", "false"));               // Use Ahead-of-Time compilation
     public static final boolean SHOW_PERF_INTERACTIVE = Boolean.parseBoolean(System.getProperty("llama.ShowPerfInteractive", "true")); // Show performance metrics in interactive mode
 
-
-    /**
-     * Loads the language model based on the given options.
-     * <p>
-     * If Ahead-of-Time (AOT) mode is enabled, attempts to use a pre-loaded compiled model. Otherwise, loads the model from the specified path using the model loader.
-     * </p>
-     *
-     * @param options
-     *         the parsed CLI options containing model path and max token limit
-     * @return the loaded {@link Model} instance
-     * @throws IOException
-     *         if the model fails to load
-     * @throws IllegalStateException
-     *         if AOT loading is enabled but the preloaded model is unavailable
-     */
-    private static Model loadModel(Options options) throws IOException {
-        if (USE_AOT) {
-            Model model = AOT.tryUsePreLoaded(options.modelPath(), options.maxTokens());
-            if (model == null) {
-                throw new IllegalStateException("Failed to load precompiled AOT model.");
-            }
-            return model;
-        }
-        return ModelLoader.loadModel(options.modelPath(), options.maxTokens(), true);
-    }
-
-    private static Sampler createSampler(Model model, Options options) {
-        return selectSampler(model.configuration().vocabularySize(), options.temperature(), options.topp(), options.seed());
-    }
-
     private static void runSingleInstruction(Model model, Sampler sampler, Options options) {
         String response = model.runInstructOnce(sampler, options);
         System.out.println(response);

src/main/java/org/beehive/gpullama3/model/loader/ModelLoader.java

@@ -1,6 +1,7 @@
 package org.beehive.gpullama3.model.loader;
 
 import org.beehive.gpullama3.Options;
+import org.beehive.gpullama3.aot.AOT;
 import org.beehive.gpullama3.core.model.GGMLType;
 import org.beehive.gpullama3.core.model.GGUF;
 import org.beehive.gpullama3.core.model.tensor.ArrayFloatTensor;
@@ -35,6 +36,8 @@
 
 public abstract class ModelLoader {
 
+    public static final boolean USE_AOT = Boolean.parseBoolean(System.getProperty("llama.AOT", "false"));               // Use Ahead-of-Time compilation
+
     protected FileChannel fileChannel;
     protected GGUF gguf;
     protected int contextLength;
@@ -74,6 +77,31 @@ private static ModelType detectModelType(Map<String, Object> metadata) {
         return ModelType.UNKNOWN;
     }
 
+    /**
+     //     * Loads the language model based on the given options.
+     //     * <p>
+     //     * If Ahead-of-Time (AOT) mode is enabled, attempts to use a pre-loaded compiled model. Otherwise, loads the model from the specified path using the model loader.
+     //     * </p>
+     //     *
+     //     * @param options
+     //     *         the parsed CLI options containing model path and max token limit
+     //     * @return the loaded {@link Model} instance
+     //     * @throws IOException
+     //     *         if the model fails to load
+     //     * @throws IllegalStateException
+     //     *         if AOT loading is enabled but the preloaded model is unavailable
+     //     */
+    public static Model loadModel(Options options) throws IOException {
+            if (USE_AOT) {
+                Model model = AOT.tryUsePreLoaded(options.modelPath(), options.maxTokens());
+                if (model == null) {
+                    throw new IllegalStateException("Failed to load precompiled AOT model.");
+                }
+                return model;
+            }
+            return ModelLoader.loadModel(options.modelPath(), options.maxTokens(), true);
+        }
+
     public static Model loadModel(Path ggufPath, int contextLength, boolean loadWeights) throws IOException {
         // initial load of metadata from gguf file
         GGUF gguf = GGUF.loadModel(ggufPath);