From b1883a066c2f2eecbc1c25015e037e23907f1613 Mon Sep 17 00:00:00 2001
From: Hansong Zhang <hsz@meta.com>
Date: Tue, 13 May 2025 15:54:37 -0700
Subject: [PATCH 1/2] [LlamaDemo] skip some tokens

---
 .../com/example/executorchllamademo/MainActivity.java    | 3 +++
 .../com/example/executorchllamademo/PromptFormat.java    | 9 +++++++++
 2 files changed, 12 insertions(+)

diff --git a/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java b/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java
index 37268202b69..b177dac69d9 100644
--- a/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java
+++ b/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java
@@ -90,6 +90,9 @@ public void onResult(String result) {
     if (result.equals(PromptFormat.getStopToken(mCurrentSettingsFields.getModelType()))) {
       return;
     }
+    if (PromptFormat.shouldSkipToken(mCurrentSettingsFields.getModelType(), result)) {
+      return;
+    }
     if (result.equals("\n\n") || result.equals("\n")) {
       if (!mResultMessage.getText().isEmpty()) {
         mResultMessage.appendText(result);
diff --git a/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/PromptFormat.java b/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/PromptFormat.java
index 5f8ecdd8042..1348cdad169 100644
--- a/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/PromptFormat.java
+++ b/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/PromptFormat.java
@@ -99,6 +99,15 @@ public static String getThinkingModeToken(ModelType modelType, boolean thinkingM
     }
   }
 
+  public static boolean shouldSkipToken(ModelType modelType, String token) {
+    switch (modelType) {
+      case QWEN_3:
+        return token.equals("<|im_end|>");
+      default:
+        return false;
+    }
+  }
+
   public static String getLlavaPresetPrompt() {
     return "A chat between a curious human and an artificial intelligence assistant. The assistant"
         + " gives helpful, detailed, and polite answers to the human's questions. USER: ";

From 311247812debf464d64ec79bcbea701007b68c49 Mon Sep 17 00:00:00 2001
From: Hansong Zhang <hsz@meta.com>
Date: Tue, 13 May 2025 16:20:22 -0700
Subject: [PATCH 2/2] replace tokens

---
 .../example/executorchllamademo/MainActivity.java   |  4 +---
 .../example/executorchllamademo/PromptFormat.java   | 13 ++++++++++---
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java b/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java
index b177dac69d9..b26031d89a6 100644
--- a/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java
+++ b/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java
@@ -90,9 +90,7 @@ public void onResult(String result) {
     if (result.equals(PromptFormat.getStopToken(mCurrentSettingsFields.getModelType()))) {
       return;
     }
-    if (PromptFormat.shouldSkipToken(mCurrentSettingsFields.getModelType(), result)) {
-      return;
-    }
+    result = PromptFormat.replaceSpecialToken(mCurrentSettingsFields.getModelType(), result);
     if (result.equals("\n\n") || result.equals("\n")) {
       if (!mResultMessage.getText().isEmpty()) {
         mResultMessage.appendText(result);
diff --git a/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/PromptFormat.java b/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/PromptFormat.java
index 1348cdad169..524ad7cbf6d 100644
--- a/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/PromptFormat.java
+++ b/examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/PromptFormat.java
@@ -99,12 +99,19 @@ public static String getThinkingModeToken(ModelType modelType, boolean thinkingM
     }
   }
 
-  public static boolean shouldSkipToken(ModelType modelType, String token) {
+  public static String replaceSpecialToken(ModelType modelType, String token) {
     switch (modelType) {
       case QWEN_3:
-        return token.equals("<|im_end|>");
+        switch (token) {
+          case "<|im_end|>":
+            return "";
+          case "<think>":
+            return "Thinking...\n";
+          case "</think>":
+            return "\nDone thinking";
+        }
       default:
-        return false;
+        return token;
     }
   }