From dd09094a227505e87d289a11c62919cbea8cebaf Mon Sep 17 00:00:00 2001
From: Xuan Son Nguyen <son@huggingface.co>
Date: Wed, 11 Dec 2024 19:28:13 +0100
Subject: [PATCH] add tok/s info

---
 examples/server/webui/index.html  | 44 ++++++++++++++++++++++---------
 examples/server/webui/src/main.js | 25 ++++++++++++++++--
 2 files changed, 55 insertions(+), 14 deletions(-)

diff --git a/examples/server/webui/index.html b/examples/server/webui/index.html
index c0bb5e634..84f86ef91 100644
--- a/examples/server/webui/index.html
+++ b/examples/server/webui/index.html
@@ -121,6 +121,7 @@
           </div>
           <div v-for="msg in messages" class="group">
             <message-bubble
+              :config="config"
               :msg="msg"
               :key="msg.id"
               :is-generating="isGenerating"
@@ -128,15 +129,16 @@
               :regenerate-msg="regenerateMsg"></message-bubble>
           </div>
 
-
-
           <!-- pending (ongoing) assistant message -->
-          <div id="pending-msg" class="chat chat-start">
-            <div v-if="pendingMsg" class="chat-bubble markdown chat-bubble-base-300">
-              <span v-if="!pendingMsg.content" class="loading loading-dots loading-md"></span>
-              <vue-markdown v-else :source="pendingMsg.content" />
-            </div>
-          </div>
+          <div id="pending-msg"></div>
+          <message-bubble
+            v-if="pendingMsg"
+            :config="config"
+            :msg="pendingMsg"
+            :key="pendingMsg.id"
+            :is-generating="isGenerating"
+            :edit-user-msg-and-regenerate="() => {}"
+            :regenerate-msg="() => {}"></message-bubble>
         </div>
 
         <!-- chat input -->
@@ -199,7 +201,7 @@
             <summary class="collapse-title font-bold">Advanced config</summary>
             <div class="collapse-content">
               <div class="flex flex-row items-center mb-2">
-                <input type="checkbox" class="checkbox" v-model="config.show_tokens_per_second" />
+                <input type="checkbox" class="checkbox" v-model="config.showTokensPerSecond" />
                 <span class="ml-4">Show tokens per second</span>
               </div>
               <label class="form-control mb-2">
@@ -243,12 +245,30 @@
           <button class="btn btn-ghost mt-2 mr-2" @click="editingContent = null">Cancel</button>
           <button class="btn mt-2" @click="editMsg()">Submit</button>
         </template>
-        <!-- render message as markdown -->
-        <vue-markdown v-else :source="msg.content" />
+        <template v-else>
+          <!-- show loading dots for pending message -->
+          <span v-if="msg.content === null" class="loading loading-dots loading-md"></span>
+          <!-- render message as markdown -->
+          <vue-markdown v-else :source="msg.content"></vue-markdown>
+          <!-- render timings if enabled -->
+          <div class="dropdown dropdown-hover dropdown-top mt-2" v-if="timings && config.showTokensPerSecond">
+            <div tabindex="0" role="button" class="cursor-pointer font-semibold text-sm opacity-70">Speed: {{ timings.predicted_per_second.toFixed(1) }} t/s</div>
+            <div class="dropdown-content bg-base-100 z-10 w-64 p-2 shadow mt-4">
+              <b>Prompt</b><br/>
+              - Tokens: {{ timings.prompt_n }}<br/>
+              - Time: {{ timings.prompt_ms }} ms<br/>
+              - Speed: {{ timings.prompt_per_second.toFixed(1) }} t/s<br/>
+              <b>Generation</b><br/>
+              - Tokens: {{ timings.predicted_n }}<br/>
+              - Time: {{ timings.predicted_ms }} ms<br/>
+              - Speed: {{ timings.predicted_per_second.toFixed(1) }} t/s<br/>
+            </div>
+          </div>
+        </template>
       </div>
     </div>
     <!-- actions for each message -->
-    <div :class="{'text-right': msg.role === 'user'}" class="mx-4 mt-2 mb-2">
+    <div :class="{'text-right': msg.role === 'user', 'opacity-0': isGenerating}" class="mx-4 mt-2 mb-2">
       <!-- user message -->
       <button v-if="msg.role === 'user'" class="badge btn-mini show-on-hover" @click="editingContent = msg.content" :disabled="isGenerating">
         ✍️ Edit
diff --git a/examples/server/webui/src/main.js b/examples/server/webui/src/main.js
index e1accb44a..53adfb159 100644
--- a/examples/server/webui/src/main.js
+++ b/examples/server/webui/src/main.js
@@ -19,6 +19,7 @@ const CONFIG_DEFAULT = {
   // Note: in order not to introduce breaking changes, please keep the same data type (number, string, etc) if you want to change the default value. Do not use null or undefined for default value.
   apiKey: '',
   systemMessage: 'You are a helpful assistant.',
+  showTokensPerSecond: false,
   // make sure these default values are in sync with `common.h`
   samplers: 'dkypmxt',
   temperature: 0.8,
@@ -39,7 +40,6 @@ const CONFIG_DEFAULT = {
   dry_allowed_length: 2,
   dry_penalty_last_n: -1,
   max_tokens: -1,
-  show_tokens_per_second: false,
   custom: '', // custom json-stringified object
 };
 const CONFIG_INFO = {
@@ -112,6 +112,7 @@ const MessageBubble = defineComponent({
   },
   template: document.getElementById('message-bubble').innerHTML,
   props: {
+    config: Object,
     msg: Object,
     isGenerating: Boolean,
     editUserMsgAndRegenerate: Function,
@@ -122,6 +123,16 @@ const MessageBubble = defineComponent({
       editingContent: null,
     };
   },
+  computed: {
+    timings() {
+      if (!this.msg.timings) return null;
+      return {
+        ...this.msg.timings,
+        prompt_per_second: this.msg.timings.prompt_n / (this.msg.timings.prompt_ms / 1000),
+        predicted_per_second: this.msg.timings.predicted_n / (this.msg.timings.predicted_ms / 1000),
+      };
+    }
+  },
   methods: {
     copyMsg() {
       copyStr(this.msg.content);
@@ -381,7 +392,7 @@ const mainApp = createApp({
           dry_allowed_length: this.config.dry_allowed_length,
           dry_penalty_last_n: this.config.dry_penalty_last_n,
           max_tokens: this.config.max_tokens,
-          timings_per_token: !!this.config.show_tokens_per_second,
+          timings_per_token: !!this.config.showTokensPerSecond,
           ...(this.config.custom.length ? JSON.parse(this.config.custom) : {}),
         };
         const chunks = sendSSEPostRequest(`${BASE_URL}/v1/chat/completions`, {
@@ -404,6 +415,16 @@ const mainApp = createApp({
               content: lastContent + addedContent,
             };
           }
+          const timings = chunk.timings;
+          if (timings) {
+            // only extract what's really needed, to save some space
+            this.pendingMsg.timings = {
+              prompt_n: timings.prompt_n,
+              prompt_ms: timings.prompt_ms,
+              predicted_n: timings.predicted_n,
+              predicted_ms: timings.predicted_ms,
+            };
+          }
         }
 
         StorageUtils.appendMsg(currConvId, this.pendingMsg);