Decouple LastRunMetrics class from Llama and reuse it for Mistral

Author: orionpapadakis

src/main/java/com/example/auxiliary/LastRunMetrics.java

@@ -0,0 +1,33 @@
+package com.example.auxiliary;
+
+/**
+ * Record to store metrics from the last model run.
+ * @param totalTokens The total number of tokens processed
+ * @param totalSeconds The total time in seconds
+ */
+public record LastRunMetrics(int totalTokens, double totalSeconds) {
+    /**
+     * Singleton instance to store the latest metrics
+     */
+    private static LastRunMetrics latestMetrics;
+
+    /**
+     * Sets the metrics for the latest run
+     *
+     * @param tokens The total number of tokens processed
+     * @param seconds The total time in seconds
+     */
+    public static void setMetrics(int tokens, double seconds) {
+        latestMetrics = new LastRunMetrics(tokens, seconds);
+    }
+
+    /**
+     * Prints the metrics from the latest run to stderr
+     */
+    public static void printMetrics() {
+        if (latestMetrics != null) {
+            double tokensPerSecond = latestMetrics.totalTokens() / latestMetrics.totalSeconds();
+            System.err.printf("\n\nachieved tok/s: %.2f. Tokens: %d, seconds: %.2f\n", tokensPerSecond, latestMetrics.totalTokens(), latestMetrics.totalSeconds());
+        }
+    }
+}

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

@@ -1,7 +1,8 @@
 package com.example.model.llama;
 
-import com.example.auxiliary.Parallel;
+import com.example.auxiliary.LastRunMetrics;
 import com.example.auxiliary.format.LlamaChatFormat;
+import com.example.auxiliary.Parallel;
 import com.example.core.model.tensor.FloatTensor;
 import com.example.model.Configuration;
 import com.example.inference.sampler.Sampler;
@@ -480,7 +481,7 @@ public void runInteractive(Sampler sampler, Options options) {
 
                 // Optionally print performance metrics after each response
                 if (SHOW_PERF_INTERACTIVE) {
-                    Llama.LastRunMetrics.printMetrics();
+                    LastRunMetrics.printMetrics();
                 }
             }
         } finally {
@@ -538,44 +539,12 @@ public void runInstructOnce(Sampler sampler, Options options) {
             System.out.println(responseText);
         }
 
-        Llama.LastRunMetrics.printMetrics();
+        LastRunMetrics.printMetrics();
 
         if (tornadoVMPlan != null) {
             tornadoVMPlan.freeTornadoExecutionPlan();
         }
     }
 
-    /**
-     * Record to store metrics from the last model run.
-     * @param totalTokens The total number of tokens processed
-     * @param totalSeconds The total time in seconds
-     */
-    public record LastRunMetrics(int totalTokens, double totalSeconds) {
-        /**
-         * Singleton instance to store the latest metrics
-         */
-        private static LastRunMetrics latestMetrics;
-
-        /**
-         * Sets the metrics for the latest run
-         *
-         * @param tokens The total number of tokens processed
-         * @param seconds The total time in seconds
-         */
-        public static void setMetrics(int tokens, double seconds) {
-            latestMetrics = new LastRunMetrics(tokens, seconds);
-        }
-
-        /**
-         * Prints the metrics from the latest run to stderr
-         */
-        public static void printMetrics() {
-            if (latestMetrics != null) {
-                double tokensPerSecond = latestMetrics.totalTokens() / latestMetrics.totalSeconds();
-                System.err.printf("\n\nachieved tok/s: %.2f. Tokens: %d, seconds: %.2f\n", tokensPerSecond, latestMetrics.totalTokens(), latestMetrics.totalSeconds());
-            }
-        }
-    }
-
 }
 

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

@@ -1,6 +1,7 @@
 package com.example.model.mistral;
 
 import com.example.auxiliary.Parallel;
+import com.example.auxiliary.LastRunMetrics;
 import com.example.auxiliary.format.MistralChatFormat;
 import com.example.core.model.tensor.FloatTensor;
 import com.example.model.Configuration;
@@ -323,5 +324,7 @@ public void runInstructOnce(Sampler sampler, Options options) {
             String responseText = tokenizer.decode(responseTokens);
             System.out.println(responseText);
         }
+
+        LastRunMetrics.printMetrics();
     }
 }