[Android] Use same stats as llm::Stats

Differential Revision: D73207250

Pull Request resolved: pytorch#10247

Author: kirklandsign

examples/demo-apps/android/LlamaDemo/app/build.gradle.kts

@@ -60,6 +60,7 @@ dependencies {
   implementation(files("libs/executorch.aar"))
   implementation("com.google.android.material:material:1.12.0")
   implementation("androidx.activity:activity:1.9.0")
+  implementation("org.json:json:20250107")
   testImplementation("junit:junit:4.13.2")
   androidTestImplementation("androidx.test.ext:junit:1.1.5")
   androidTestImplementation("androidx.test.espresso:espresso-core:3.5.1")

examples/demo-apps/android/LlamaDemo/app/src/androidTest/java/com/example/executorchllamademo/PerfTest.java

@@ -18,6 +18,8 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import org.json.JSONException;
+import org.json.JSONObject;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.pytorch.executorch.extension.llm.LlmCallback;
@@ -64,8 +66,16 @@ public void onResult(String result) {
   }
 
   @Override
-  public void onStats(float tps) {
-    tokensPerSecond.add(tps);
+  public void onStats(String result) {
+    try {
+      JSONObject jsonObject = new JSONObject(result);
+      int numGeneratedTokens = jsonObject.getInt("generated_tokens");
+      int inferenceEndMs = jsonObject.getInt("inference_end_ms");
+      int promptEvalEndMs = jsonObject.getInt("prompt_eval_end_ms");
+      float tps = (float) numGeneratedTokens / (inferenceEndMs - promptEvalEndMs) * 1000;
+      tokensPerSecond.add(tps);
+    } catch (JSONException e) {
+    }
   }
 
   private void report(final String metric, final Float value) {

examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java

@@ -49,6 +49,8 @@
 import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
+import org.json.JSONException;
+import org.json.JSONObject;
 import org.pytorch.executorch.extension.llm.LlmCallback;
 import org.pytorch.executorch.extension.llm.LlmModule;
 
@@ -97,10 +99,20 @@ public void onResult(String result) {
   }
 
   @Override
-  public void onStats(float tps) {
+  public void onStats(String stats) {
     runOnUiThread(
         () -> {
           if (mResultMessage != null) {
+            float tps = 0;
+            try {
+              JSONObject jsonObject = new JSONObject(stats);
+              int numGeneratedTokens = jsonObject.getInt("generated_tokens");
+              int inferenceEndMs = jsonObject.getInt("inference_end_ms");
+              int promptEvalEndMs = jsonObject.getInt("prompt_eval_end_ms");
+              tps = (float) numGeneratedTokens / (inferenceEndMs - promptEvalEndMs) * 1000;
+            } catch (JSONException e) {
+              Log.e("LLM", "Error parsing JSON: " + e.getMessage());
+            }
             mResultMessage.setTokensPerSecond(tps);
             mMessageAdapter.notifyDataSetChanged();
           }

examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ModelRunner.java

@@ -13,6 +13,8 @@
 import android.os.Looper;
 import android.os.Message;
 import androidx.annotation.NonNull;
+import org.json.JSONException;
+import org.json.JSONObject;
 import org.pytorch.executorch.extension.llm.LlmCallback;
 import org.pytorch.executorch.extension.llm.LlmModule;
 
@@ -69,7 +71,16 @@ public void onResult(String result) {
   }
 
   @Override
-  public void onStats(float tps) {
+  public void onStats(String stats) {
+    float tps = 0;
+    try {
+      JSONObject jsonObject = new JSONObject(stats);
+      int numGeneratedTokens = jsonObject.getInt("generated_tokens");
+      int inferenceEndMs = jsonObject.getInt("inference_end_ms");
+      int promptEvalEndMs = jsonObject.getInt("prompt_eval_end_ms");
+      tps = (float) numGeneratedTokens / (inferenceEndMs - promptEvalEndMs) * 1000;
+    } catch (JSONException e) {
+    }
     mCallback.onStats("tokens/second: " + tps);
   }
 }

examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/ModelRunnerCallback.java

@@ -18,7 +18,7 @@ public interface ModelRunnerCallback {
 
   void onTokenGenerated(String token);
 
-  void onStats(String token);
+  void onStats(String stats);
 
   void onGenerationStopped();
 }

extension/android/executorch_android/build.gradle

@@ -47,6 +47,7 @@ dependencies {
     androidTestImplementation 'androidx.test.ext:junit:1.1.5'
     androidTestImplementation 'androidx.test:rules:1.2.0'
     androidTestImplementation 'commons-io:commons-io:2.4'
+    androidTestImplementation 'org.json:json:20250107'
 }
 
 import com.vanniktech.maven.publish.SonatypeHost

extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/LlmModuleInstrumentationTest.java

@@ -34,6 +34,8 @@
 import org.apache.commons.io.FileUtils;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.InstrumentationRegistry;
+import org.json.JSONException;
+import org.json.JSONObject;
 import org.pytorch.executorch.extension.llm.LlmCallback;
 import org.pytorch.executorch.extension.llm.LlmModule;
 
@@ -94,8 +96,17 @@ public void onResult(String result) {
             }
 
             @Override
-            public void onStats(float tps) {
-                LlmModuleInstrumentationTest.this.onStats(tps);
+            public void onStats(String stats) {
+                float tps = 0;
+                try {
+                    JSONObject jsonObject = new JSONObject(stats);
+                    int numGeneratedTokens = jsonObject.getInt("generated_tokens");
+                    int inferenceEndMs = jsonObject.getInt("inference_end_ms");
+                    int promptEvalEndMs = jsonObject.getInt("prompt_eval_end_ms");
+                    tps = (float) numGeneratedTokens / (inferenceEndMs - promptEvalEndMs) * 1000;
+                    LlmModuleInstrumentationTest.this.onStats(tps);
+                } catch (JSONException e) {
+                }
             }
         });
 

extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmCallback.java

@@ -31,8 +31,22 @@ public interface LlmCallback {
   /**
    * Called when the statistics for the generate() is available.
    *
+   * Note: This is a deprecated API and will be removed in the future. Please use onStats(String stats)
+   *
    * @param tps Tokens/second for generated tokens.
    */
+  @Deprecated
+  @DoNotStrip
+  default public void onStats(float tps) {}
+
+  /**
+   * Called when the statistics for the generate() is available.
+   *
+   * The result will be a JSON string. See extension/llm/stats.h for the field
+   * definitions.
+   *
+   * @param stats JSON string containing the statistics for the generate()
+   */
   @DoNotStrip
-  public void onStats(float tps);
+  default public void onStats(String stats) {}
 }

extension/android/jni/jni_layer_llama.cpp

@@ -100,14 +100,20 @@ class ExecuTorchLlmCallbackJni
 
   void onStats(const llm::Stats& result) const {
     static auto cls = ExecuTorchLlmCallbackJni::javaClassStatic();
-    static const auto method = cls->getMethod<void(jfloat)>("onStats");
+    static const auto tps_method = cls->getMethod<void(jfloat)>("onStats");
     double eval_time =
         (double)(result.inference_end_ms - result.prompt_eval_end_ms);
 
     float tps = result.num_generated_tokens / eval_time *
         result.SCALING_FACTOR_UNITS_PER_SECOND;
-
-    method(self(), tps);
+    tps_method(self(), tps);
+
+    static const auto on_stats_method =
+        cls->getMethod<void(facebook::jni::local_ref<jstring>)>("onStats");
+    on_stats_method(
+        self(),
+        facebook::jni::make_jstring(
+            executorch::extension::llm::stats_to_json_string(result)));
   }
 };
 

extension/benchmark/android/benchmark/app/build.gradle.kts

@@ -39,6 +39,7 @@ dependencies {
   implementation("com.facebook.soloader:soloader:0.10.5")
   implementation("com.facebook.fbjni:fbjni:0.5.1")
   implementation("com.google.code.gson:gson:2.8.6")
+  implementation("org.json:json:20250107")
   testImplementation("junit:junit:4.13.2")
   androidTestImplementation("androidx.test.ext:junit:1.2.1")
   androidTestImplementation("androidx.test.espresso:espresso-core:3.6.1")