json

Author: kirklandsign

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")

extension/benchmark/android/benchmark/app/src/main/java/org/pytorch/minibench/LlmBenchmarkActivity.java

@@ -23,6 +23,8 @@
 import java.util.List;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import org.json.JSONException;
+import org.json.JSONObject;
 
 public class LlmBenchmarkActivity extends Activity implements ModelRunnerCallback {
   ModelRunner mModelRunner;
@@ -80,8 +82,17 @@ public void onTokenGenerated(String token) {}
 
   @Override
   public void onStats(String result) {
-    // TODO: Add tps
-    // mStatsInfo.tokens = stats;
+    float tps = 0;
+    try {
+      JSONObject jsonObject = new JSONObject(stats);
+      int numGeneratedTokens = jsonObject.getInt("num_generated_tokens");
+      int inferenceEndMs = jsonObject.getInt("inference_end_ms");
+      int promptEvalEndMs = jsonObject.getInt("prompt_eval_end_ms");
+      tps = (float) numGeneratedTokens / (inferenceEndMs - promptEvalEndMs) * 1000;
+      mStatsInfo.tokens = String.valueOf(tps);
+    } catch (JSONException e) {
+      Log.e("LLM", "Error parsing JSON: " + e.getMessage());
+    }
   }
 
   @Override