Try with any pte

Author: kirklandsign

extension/android/benchmark/android-llm-device-farm-test-spec.yml

@@ -80,6 +80,12 @@ phases:
         --es "model_dir" "/data/local/tmp/llama" \
         --es "tokenizer_path" "/data/local/tmp/llama/tokenizer.bin"
 
+      - echo "Run generic benchmark"
+      - |
+        adb -s $DEVICEFARM_DEVICE_UDID shell am start -W -n org.pytorch.minibench/.BenchmarkActivity \
+        --es "model_dir" "/data/local/tmp/llama"
+
+
   post_test:
     commands:
       - echo "Gather LLM benchmark results"

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

@@ -11,22 +11,30 @@
 import android.app.Activity;
 import android.content.Intent;
 import android.os.Bundle;
+import java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
+import java.util.Arrays;
 import org.pytorch.executorch.Module;
 
 public class BenchmarkActivity extends Activity {
   @Override
   protected void onCreate(Bundle savedInstanceState) {
     super.onCreate(savedInstanceState);
     Intent intent = getIntent();
-    String modelPath = intent.getStringExtra("model_path");
+    File modelDir = new File(intent.getStringExtra("model_dir"));
+    File model =
+        Arrays.stream(modelDir.listFiles())
+            .filter(file -> file.getName().endsWith(".pte"))
+            .findFirst()
+            .get();
+
     int numIter = intent.getIntExtra("num_iter", 10);
 
     // TODO: Format the string with a parsable format
     StringBuilder resultText = new StringBuilder();
 
-    Module module = Module.load(modelPath);
+    Module module = Module.load(model.getPath());
     for (int i = 0; i < numIter; i++) {
       long start = System.currentTimeMillis();
       module.forward();