Polish UX and display t/s (#2286)

Summary: Pull Request resolved: #2286 Differential Revision: D54612043
pytorch · Mar 7, 2024 · 0a95d0c · 0a95d0c
1 parent cf929d0
commit 0a95d0c
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 0 deletions.
diff --git a/...emo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchdemo/MainActivity.java b/...emo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchdemo/MainActivity.java
@@ -28,10 +28,14 @@ public class MainActivity extends Activity implements Runnable, LlamaCallback {
   private LlamaModule mModule = null;
   private Message mResultMessage = null;
 
+  private int mNumTokens = 0;
+  private long mRunStartTime = 0;
+
   @Override
   public void onResult(String result) {
     System.out.println("onResult: " + result);
     mResultMessage.appendText(result);
+    mNumTokens++;
     run();
   }
 
@@ -87,7 +91,23 @@ protected void onCreate(Bundle savedInstanceState) {
               new Runnable() {
                 @Override
                 public void run() {
+                  runOnUiThread(
+                      new Runnable() {
+                        @Override
+                        public void run() {
+                          onModelRunStarted();
+                        }
+                      });
+
                   mModule.generate(prompt, MainActivity.this);
+
+                  runOnUiThread(
+                      new Runnable() {
+                        @Override
+                        public void run() {
+                          onModelRunStopped();
+                        }
+                      });
                 }
               };
           new Thread(runnable).start();
@@ -107,6 +127,25 @@ public void run() {
         });
 
     setLocalModel("/data/local/tmp/stories110M.pte", "/data/local/tmp/tokenizer.bin");
+    onModelRunStopped();
+  }
+
+  private void onModelRunStarted() {
+    mSendButton.setEnabled(false);
+    mStopButton.setEnabled(true);
+    mRunStartTime = System.currentTimeMillis();
+  }
+
+  private void onModelRunStopped() {
+    long runDuration = System.currentTimeMillis() - mRunStartTime;
+    if (mResultMessage != null) {
+      mResultMessage.setTokensPerSecond(1.0f * mNumTokens / (runDuration / 1000.0f));
+    }
+    mSendButton.setEnabled(true);
+    mStopButton.setEnabled(false);
+    mNumTokens = 0;
+    mRunStartTime = 0;
+    mMessageAdapter.notifyDataSetChanged();
   }
 
   @Override

diff --git a/...les/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchdemo/Message.java b/...les/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchdemo/Message.java
@@ -11,6 +11,7 @@
 public class Message {
   private String text;
   private boolean isSent;
+  private float tokensPerSecond;
 
   public Message(String text, boolean isSent) {
     this.text = text;
@@ -28,4 +29,12 @@ public void appendText(String text) {
   public boolean getIsSent() {
     return isSent;
   }
+
+  public void setTokensPerSecond(float tokensPerSecond) {
+    this.tokensPerSecond = tokensPerSecond;
+  }
+
+  public float getTokensPerSecond() {
+    return tokensPerSecond;
+  }
 }
diff --git a/...o-apps/android/LlamaDemo/app/src/main/java/com/example/executorchdemo/MessageAdapter.java b/...o-apps/android/LlamaDemo/app/src/main/java/com/example/executorchdemo/MessageAdapter.java
@@ -28,6 +28,12 @@ public View getView(int position, View convertView, ViewGroup parent) {
         LayoutInflater.from(getContext()).inflate(layoutIdForListItem, parent, false);
     TextView messageTextView = listItemView.findViewById(R.id.message_text);
     messageTextView.setText(currentMessage.getText());
+
+    if (currentMessage.getTokensPerSecond() > 0) {
+      TextView tokensView = listItemView.findViewById(R.id.tokens_per_second);
+      tokensView.setText("" + currentMessage.getTokensPerSecond() + " t/s");
+    }
+
     return listItemView;
   }
 }
diff --git a/examples/demo-apps/android/LlamaDemo/app/src/main/res/layout/received_message.xml b/examples/demo-apps/android/LlamaDemo/app/src/main/res/layout/received_message.xml
@@ -28,4 +28,15 @@
         android:textSize="18dp"
         android:text="Generated text"
         />
+
+
+    <TextView
+        android:id="@+id/tokens_per_second"
+        android:layout_marginLeft="15dp"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:layout_below="@+id/message_text"
+        android:paddingBottom="4dp"
+        android:text=""/>
+
 </RelativeLayout>