Refactor LlamaApp.java to improve model loading and sampler creation

Author: mikepapadim

src/main/java/com/example/LlamaApp.java

@@ -5,8 +5,8 @@
 import com.example.inference.sampler.CategoricalSampler;
 import com.example.inference.sampler.Sampler;
 import com.example.inference.sampler.ToppSampler;
-import com.example.model.Model;
 import com.example.loader.weights.ModelLoader;
+import com.example.model.Model;
 import com.example.tornadovm.FloatArrayUtils;
 import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
 
@@ -106,16 +106,48 @@ static Sampler selectSampler(int vocabularySize, float temperature, float topp,
         return sampler;
     }
 
-    public static void main(String[] args) throws IOException {
-        Options options = Options.parseOptions(args);
-        Model model;
+    /**
+     * 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 = AOT.tryUsePreLoaded(options.modelPath(), options.maxTokens());
-        } else {
-            model = ModelLoader.loadModel(options.modelPath(), options.maxTokens(), true);
+            Model model = AOT.tryUsePreLoaded(options.modelPath(), options.maxTokens());
+            if (model == null) {
+                throw new IllegalStateException("Failed to load precompiled AOT model.");
+            }
+            return model;
         }
-        assert model != null;
-        Sampler sampler = selectSampler(model.configuration().vocabularySize(), options.temperature(), options.topp(), options.seed());
+        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());
+    }
+
+    /**
+     * Entry point for running the LLaMA-based model with provided command-line arguments.
+     *
+     * <p>Initializes model options, loads the appropriate model (either AOT or on-demand),
+     * configures the sampler, and runs either in interactive or single-instruction mode
+     * based on the input options.</p>
+     *
+     * @param args command-line arguments used to configure model path, temperature, seed, etc.
+     * @throws IOException if model loading or file operations fail.
+     */
+    public static void main(String[] args) throws IOException {
+        Options options = Options.parseOptions(args);
+        Model model = loadModel(options);
+        Sampler sampler = createSampler(model, options);
+
         if (options.interactive()) {
             model.runInteractive(sampler, options);
         } else {