Use same as llm:Stats

Author: kirklandsign

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
   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
+  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)));
   }
 };