Add backend parameter to ModelUtil

Author: cmodi-meta

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

@@ -125,8 +125,7 @@ private void setLocalModel(String modelPath, String tokenizerPath, float tempera
     long runStartTime = System.currentTimeMillis();
     mModule =
         new LlamaModule(
-            //ModelUtils.getModelCategory(mCurrentSettingsFields.getModelType()),
-                3, //TODO: Modify this based on JNI change for how to select MTK backend
+            ModelUtils.getModelCategory(mCurrentSettingsFields.getModelType(), mCurrentSettingsFields.getBackendType()),
             modelPath,
             tokenizerPath,
             temperature);
@@ -175,6 +174,10 @@ private void setLocalModel(String modelPath, String tokenizerPath, float tempera
             + modelPath
             + "\nTokenizer path: "
             + tokenizerPath
+            + "\nBackend: "
+            + mCurrentSettingsFields.getBackendType().toString()
+            + "\nModelType: "
+            + ModelUtils.getModelCategory(mCurrentSettingsFields.getModelType(), mCurrentSettingsFields.getBackendType())
             + "\nTemperature: "
             + temperature
             + "\nModel loaded time: "
@@ -692,7 +695,7 @@ private void onModelRunStopped() {
           addSelectedImagesToChatThread(mSelectedImageUri);
           String finalPrompt;
           String rawPrompt = mEditTextMessage.getText().toString();
-          if (ModelUtils.getModelCategory(mCurrentSettingsFields.getModelType())
+          if (ModelUtils.getModelCategory(mCurrentSettingsFields.getModelType(), mCurrentSettingsFields.getBackendType())
               == ModelUtils.VISION_MODEL) {
             finalPrompt = mCurrentSettingsFields.getFormattedSystemAndUserPrompt(rawPrompt);
           } else {
@@ -725,7 +728,7 @@ public void run() {
                         }
                       });
                   long generateStartTime = System.currentTimeMillis();
-                  if (ModelUtils.getModelCategory(mCurrentSettingsFields.getModelType())
+                  if (ModelUtils.getModelCategory(mCurrentSettingsFields.getModelType(), mCurrentSettingsFields.getBackendType())
                       == ModelUtils.VISION_MODEL) {
                     mModule.generateFromPos(
                         finalPrompt,

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

@@ -9,22 +9,33 @@
 package com.example.executorchllamademo;
 
 public class ModelUtils {
+  // XNNPACK or QNN
   static final int TEXT_MODEL = 1;
+
+  // XNNPACK
   static final int VISION_MODEL = 2;
   static final int VISION_MODEL_IMAGE_CHANNELS = 3;
-  //TODO: Make change here based on JNI change on how to indicate MTK backend
   static final int VISION_MODEL_SEQ_LEN = 768;
   static final int TEXT_MODEL_SEQ_LEN = 256;
 
-  public static int getModelCategory(ModelType modelType) {
-    switch (modelType) {
-      case LLAVA_1_5:
-        return VISION_MODEL;
-      case LLAMA_3:
-      case LLAMA_3_1:
-      case LLAMA_3_2:
-      default:
-        return TEXT_MODEL;
+  // MediaTek
+  static final int MEDIATEK_TEXT_MODEL = 3;
+
+  public static int getModelCategory(ModelType modelType, BackendType backendType) {
+    if (backendType.equals(BackendType.XNNPACK)) {
+      switch (modelType) {
+        case LLAVA_1_5:
+          return VISION_MODEL;
+        case LLAMA_3:
+        case LLAMA_3_1:
+        case LLAMA_3_2:
+        default:
+          return TEXT_MODEL;
+      }
+    } else if (backendType.equals(BackendType.MEDIATEK)) {
+      return MEDIATEK_TEXT_MODEL;
     }
+
+    return TEXT_MODEL; // default
   }
 }