From bb226a4147e237887afac2f440bc9833ddf46b0b Mon Sep 17 00:00:00 2001 From: rcmerci Date: Fri, 21 Mar 2025 01:38:39 +0800 Subject: [PATCH] feat(vec-search): debug ui show load-model-progress --- .../components/vector_search/sidebar.cljs | 17 +++++- .../handler/db_based/vector_search_flows.cljs | 3 ++ src/main/frontend/handler/events.cljs | 3 ++ src/main/frontend/handler/worker.cljs | 5 +- .../inference_worker/text_embedding.cljs | 12 ++++- src/main/frontend/persist_db/browser.cljs | 1 + src/main/frontend/state.cljs | 3 +- src/main/frontend/worker/db_worker.cljs | 6 ++- src/main/frontend/worker/embedding.cljs | 54 +++++++++++-------- 9 files changed, 76 insertions(+), 28 deletions(-) diff --git a/src/main/frontend/components/vector_search/sidebar.cljs b/src/main/frontend/components/vector_search/sidebar.cljs index ed915b32d8..f3a64938ff 100644 --- a/src/main/frontend/components/vector_search/sidebar.cljs +++ b/src/main/frontend/components/vector_search/sidebar.cljs @@ -17,6 +17,7 @@ ^js worker @db-browser/*worker [model-info set-model-info] (hooks/use-state nil) [vec-search-state set-vec-search-state] (hooks/use-state nil) + [load-model-progress set-load-model-progress] (hooks/use-state nil) [query-string set-query-string] (hooks/use-state nil) [result set-result] (hooks/use-state nil)] (hooks/use-effect! @@ -24,9 +25,20 @@ (c.m/run-task (m/reduce (fn [_ v] (set-vec-search-state v)) - vector-search-flows/vector-search-state-flow) + (m/ap + (m/?> vector-search-flows/infer-worker-ready-flow) + (c.m/ vector-search-flows/vector-search-state-flow))) ::update-vec-search-state :succ (constantly nil))) []) + (hooks/use-effect! + (fn [] + (c.m/run-task + (m/reduce + (fn [_ v] (set-load-model-progress v)) + vector-search-flows/load-model-progress-flow) + ::update-load-model-progress :succ (constantly nil))) + []) (hooks/use-effect! (fn [] (c.m/run-task @@ -50,7 +62,8 @@ [(hooks/use-debounced-value query-string 200)]) [:div [:b "State"] - (let [state-map (get-in vec-search-state [:repo->index-info repo])] + (let [state-map (assoc (get-in vec-search-state [:repo->index-info repo]) + :load-model-progress load-model-progress)] [:pre.select-text (with-out-str (fipp/pprint state-map {:width 10}))]) diff --git a/src/main/frontend/handler/db_based/vector_search_flows.cljs b/src/main/frontend/handler/db_based/vector_search_flows.cljs index a449949dd2..4950e78c28 100644 --- a/src/main/frontend/handler/db_based/vector_search_flows.cljs +++ b/src/main/frontend/handler/db_based/vector_search_flows.cljs @@ -15,5 +15,8 @@ (def vector-search-state-flow (m/watch (:vector-search/state @state/state))) +(def load-model-progress-flow + (m/watch (:vector-search/load-model-progress @state/state))) + (comment ((m/reduce (fn [_ x] (prn :xx x)) vector-search-state-flow) prn js/console.log)) diff --git a/src/main/frontend/handler/events.cljs b/src/main/frontend/handler/events.cljs index 299806b4e8..1cb9afe8d4 100644 --- a/src/main/frontend/handler/events.cljs +++ b/src/main/frontend/handler/events.cljs @@ -1095,6 +1095,9 @@ (defmethod handle :editor/hide-action-bar [] (shui/popup-hide! :selection-action-bar)) +(defmethod handle :vector-search/load-model-progress [[_ data]] + (state/set-state! :vector-search/load-model-progress data)) + (defn run! [] (let [chan (state/get-events-chan)] diff --git a/src/main/frontend/handler/worker.cljs b/src/main/frontend/handler/worker.cljs index 4d12a776d9..6b145cd79d 100644 --- a/src/main/frontend/handler/worker.cljs +++ b/src/main/frontend/handler/worker.cljs @@ -1,5 +1,5 @@ (ns frontend.handler.worker - "Handle messages received from the db worker" + "Handle messages received from the webworkers" (:require [cljs-bean.core :as bean] [frontend.handler.file-based.file :as file-handler] [frontend.handler.notification :as notification] @@ -52,6 +52,9 @@ (defmethod handle :capture-error [_ _worker data] (state/pub-event! [:capture-error data])) +(defmethod handle :vector-search/load-model-progress [_ _ data] + (state/pub-event! [:vector-search/load-model-progress data])) + (defmethod handle :default [_ _worker data] (prn :debug "Worker data not handled: " data)) diff --git a/src/main/frontend/inference_worker/text_embedding.cljs b/src/main/frontend/inference_worker/text_embedding.cljs index 967a23f071..071c531739 100644 --- a/src/main/frontend/inference_worker/text_embedding.cljs +++ b/src/main/frontend/inference_worker/text_embedding.cljs @@ -7,6 +7,7 @@ [frontend.inference-worker.state :as infer-worker.state] [frontend.worker-common.util :as worker-util] [lambdaisland.glogi :as log] + [logseq.common.config :as common-config] [missionary.core :as m] [promesa.core :as p])) @@ -173,13 +174,15 @@ "Xenova/jina-embeddings-v2-base-zh" {:tf-config {:dtype "fp32"} :hnsw-config {:dims 768}}}) +(def ^:private *load-model-progress (atom nil)) + (defn js (-> (:tf-config config) (assoc "device" "webgpu") - (assoc "progress_callback" #(log/info :progress %)))))] + (assoc "progress_callback" #(reset! *load-model-progress %)))))] (reset! infer-worker.state/*extractor extractor) (reset! infer-worker.state/*model-name+config [model-name config]) true))) @@ -192,6 +195,13 @@ (.setDebugLogs (.-EmscriptenFileSystemManager ^js @infer-worker.state/*hnswlib) true) (log/info :loaded :hnswlib)))) +(when-not common-config/PUBLISHING + (c.m/run-background-task + ::push-load-model-progress + (m/reduce + (fn [_ v] (worker-util/post-message :vector-search/load-model-progress v)) + (c.m/throttle 500 (m/watch *load-model-progress))))) + (comment (def repo "repo-1") (def hnsw (ensure-hnsw-index! repo)) diff --git a/src/main/frontend/persist_db/browser.cljs b/src/main/frontend/persist_db/browser.cljs index 07184d1638..0fd5502a5f 100644 --- a/src/main/frontend/persist_db/browser.cljs +++ b/src/main/frontend/persist_db/browser.cljs @@ -164,6 +164,7 @@ worker (js/Worker. (str worker-url "?electron=" (util/electron?) "&publishing=" config/publishing?)) wrapped-worker (Comlink/wrap worker) t1 (util/time-ms)] + (worker-handler/handle-message! worker wrapped-worker) (reset! state/*infer-worker wrapped-worker) (p/do! (.init wrapped-worker) diff --git a/src/main/frontend/state.cljs b/src/main/frontend/state.cljs index e1a88a6ee6..80b1526db0 100644 --- a/src/main/frontend/state.cljs +++ b/src/main/frontend/state.cljs @@ -332,7 +332,8 @@ :db/async-queries (atom {}) :db/latest-transacted-entity-uuids (atom {}) - :vector-search/state (atom {})}))) + :vector-search/state (atom {}) + :vector-search/load-model-progress (atom nil)}))) ;; Block ast state ;; =============== diff --git a/src/main/frontend/worker/db_worker.cljs b/src/main/frontend/worker/db_worker.cljs index d6a798a147..1c805b697e 100644 --- a/src/main/frontend/worker/db_worker.cljs +++ b/src/main/frontend/worker/db_worker.cljs @@ -955,7 +955,11 @@ (vec-search-cancel-indexing [this repo] - (embedding/cancel-indexing repo))) + (embedding/cancel-indexing repo)) + + (vec-search-update-index-info + [this repo] + (js/Promise. (embedding/task--update-index-info! repo)))) (defn- rename-page! [repo conn page-uuid new-name] diff --git a/src/main/frontend/worker/embedding.cljs b/src/main/frontend/worker/embedding.cljs index ed05da6e99..0133eda405 100644 --- a/src/main/frontend/worker/embedding.cljs +++ b/src/main/frontend/worker/embedding.cljs @@ -111,13 +111,25 @@ [:db/add e :logseq.property.embedding/hnsw-label-updated-at updated-at]])) e+updated-at-coll added-labels))) -(defn- task--update-index-info! - [repo ^js infer-worker indexing?*] +(defn- task--update-index-info!* + ([repo ^js infer-worker] + (m/sp + (reset-*vector-search-state! repo :index-info + (merge (:index-info @*vector-search-state) + (js->clj (c.m/clj (c.m/clj (c.m/