Fix instrumentation

Author: kirklandsign

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,9 +96,17 @@ public void onResult(String result) {
             }
 
             @Override
-            public void onStats(String result) {
-                // TODO: Calculate 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) {
+                }
             }
         });