From 9be00ea4afb33b271e2eb32f2c8ba3e25ee611dd Mon Sep 17 00:00:00 2001
From: Neko Ayaka <neko@ayaka.moe>
Date: Thu, 27 Feb 2025 19:40:14 +0800
Subject: [PATCH] fix(providers-transformers): progress support

Signed-off-by: Neko Ayaka <neko@ayaka.moe>
---
 .../playground/src/App.vue                    | 36 ++++++++++++++++++-
 .../playground/src/components/Progress.vue    | 27 ++++++++++++++
 packages/provider-transformers/src/index.ts   | 22 ++++++++----
 3 files changed, 77 insertions(+), 8 deletions(-)
 create mode 100644 packages/provider-transformers/playground/src/components/Progress.vue

diff --git a/packages/provider-transformers/playground/src/App.vue b/packages/provider-transformers/playground/src/App.vue
index d3823d0..76b9871 100644
--- a/packages/provider-transformers/playground/src/App.vue
+++ b/packages/provider-transformers/playground/src/App.vue
@@ -1,4 +1,6 @@
 <script setup lang="ts">
+import type { InitiateProgressInfo, ProgressStatusInfo } from '@proj-airi/utils-transformers/types'
+
 import { useDark, useToggle } from '@vueuse/core'
 import { embed } from '@xsai/embed'
 import { serialize } from 'superjson'
@@ -6,6 +8,7 @@ import { onMounted, ref } from 'vue'
 
 import { createTransformers } from '../../src'
 import embedWorkerURL from '../../src/worker?worker&url'
+import Progress from './components/Progress.vue'
 
 const isDark = useDark()
 const toggleDark = useToggle(isDark)
@@ -13,11 +16,33 @@ const toggleDark = useToggle(isDark)
 const modelId = ref('Xenova/all-MiniLM-L6-v2')
 const input = ref('Hello, world!')
 const results = ref<any>()
+const loadingItems = ref<(InitiateProgressInfo | ProgressStatusInfo)[]>([])
 
 const transformersProvider = createTransformers({ embedWorkerURL })
 
 onMounted(async () => {
-  await transformersProvider.loadEmbed(modelId.value)
+  await transformersProvider.loadEmbed(modelId.value, {
+    onProgress: (progress) => {
+      switch (progress.status) {
+        case 'initiate':
+          loadingItems.value.push(progress)
+          break
+
+        case 'progress':
+          loadingItems.value = loadingItems.value.map((item) => {
+            if (item.file === progress.file) {
+              return { ...item, ...progress }
+            }
+            return item
+          })
+          break
+
+        case 'done':
+          // loadingItems.value = loadingItems.value.filter(item => item.file !== progress.file)
+          break
+      }
+    },
+  })
 })
 
 async function execute() {
@@ -55,6 +80,15 @@ async function execute() {
           </label>
         </div>
       </div>
+      <div v-if="loadingItems.length > 0" class="w-[50%] flex flex-col gap-2">
+        <Progress
+          v-for="(item, index) of loadingItems"
+          :key="index"
+          :text="item.file"
+          :percentage="'progress' in item ? item.progress || 0 : 0"
+          :total="'total' in item ? item.total || 0 : 0"
+        />
+      </div>
     </div>
     <div grid grid-cols-2 gap-2>
       <div flex flex-col gap-2>
diff --git a/packages/provider-transformers/playground/src/components/Progress.vue b/packages/provider-transformers/playground/src/components/Progress.vue
new file mode 100644
index 0000000..b90c72c
--- /dev/null
+++ b/packages/provider-transformers/playground/src/components/Progress.vue
@@ -0,0 +1,27 @@
+<script setup lang="ts">
+withDefaults(defineProps<{
+  text: string
+  percentage: number
+  total?: number
+}>(), {
+  percentage: 0,
+})
+
+function formatBytes(size?: number) {
+  if (!size)
+    size = 0
+
+  const i = size === 0 ? 0 : Math.floor(Math.log(size) / Math.log(1024))
+  return +((size / 1024 ** i).toFixed(2)) * 1 + ['B', 'kB', 'MB', 'GB', 'TB'][i]
+}
+</script>
+
+<template>
+  <div
+    class="flex items-center justify-between whitespace-nowrap rounded-lg bg-violet-100 px-3 py-2 text-sm dark:bg-violet-900"
+    :style="{ width: `${percentage}%` }"
+  >
+    <span>{{ text }}</span>
+    <span>({{ percentage.toFixed(2) }}%{{ Number.isNaN(total) ? '' : ` of ${formatBytes(total)}` }})</span>
+  </div>
+</template>
diff --git a/packages/provider-transformers/src/index.ts b/packages/provider-transformers/src/index.ts
index f0eb0d4..d780194 100644
--- a/packages/provider-transformers/src/index.ts
+++ b/packages/provider-transformers/src/index.ts
@@ -15,6 +15,9 @@ export function createEmbedProvider<T extends string, T2 extends Omit<CommonRequ
 
   function loadModel(model: (string & {}) | T, options: T2) {
     return new Promise<void>((resolve, reject) => {
+      const onProgress = options.onProgress
+      delete options.onProgress
+
       try {
         const workerURL = new URL(createOptions.baseURL)
 
@@ -29,7 +32,7 @@ export function createEmbedProvider<T extends string, T2 extends Omit<CommonRequ
         reject(err)
       }
 
-      worker.onmessage = (event: MessageEvent<WorkerMessageEvent>) => {
+      worker.addEventListener('message', (event: MessageEvent<WorkerMessageEvent>) => {
         switch (event.data.type) {
           case 'error':
             reject(event.data.data.error)
@@ -42,13 +45,13 @@ export function createEmbedProvider<T extends string, T2 extends Omit<CommonRequ
 
             break
           case 'progress':
-            if (options.onProgress != null && typeof options.onProgress === 'function') {
-              options.onProgress(event.data.data.progress)
+            if (onProgress != null && typeof onProgress === 'function') {
+              onProgress(event.data.data.progress)
             }
 
             break
         }
-      }
+      })
     })
   }
 
@@ -62,13 +65,18 @@ export function createEmbedProvider<T extends string, T2 extends Omit<CommonRequ
               return
             }
 
+            worker.addEventListener('error', (event: ErrorEvent) => {
+              reject(event)
+            })
+
             let text: string = ''
-            let body: any
+            let body: LoadOptions & { input: string }
 
             try {
               body = JSON.parse(init.body.toString())
               text = body.input
               delete body.input
+              delete body.onProgress
             }
             catch (err) {
               reject(err)
@@ -78,7 +86,7 @@ export function createEmbedProvider<T extends string, T2 extends Omit<CommonRequ
             let errored = false
             let resultDone = false
 
-            worker.onmessage = (event: MessageEvent<WorkerMessageEvent>) => {
+            worker.addEventListener('message', (event: MessageEvent<WorkerMessageEvent>) => {
               switch (event.data.type) {
                 case 'error':
                   errored = true
@@ -94,7 +102,7 @@ export function createEmbedProvider<T extends string, T2 extends Omit<CommonRequ
                   resolve(new Response(encoder.encode(JSON.stringify(result))))
                   break
               }
-            }
+            })
 
             if (!errored && !resultDone)
               worker.postMessage({ type: 'extract', data: { text, options: body as any } } satisfies WorkerMessageEvent)