feat: node addon support muti input files and callback function suppo…

…rt provide progress
zcf0508 · Mar 1, 2024 · 5ecffc9 · 5ecffc9
1 parent 25d313b
commit 5ecffc9
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 18 deletions.
diff --git a/examples/addon.node/__test__/whisper.spec.js b/examples/addon.node/__test__/whisper.spec.js
@@ -10,15 +10,17 @@ const whisperAsync = promisify(whisper);
 const whisperParamsMock = {
   language: "en",
   model: path.join(__dirname, "../../../models/ggml-base.en.bin"),
-  fname_inp: path.join(__dirname, "../../../samples/jfk.wav"),
+  fname_inp: [path.join(__dirname, "../../../samples/jfk.wav")],
   use_gpu: true,
 };
 
 describe("Run whisper.node", () => {
     test("it should receive a non-empty value", async () => {
-        let result = await whisperAsync(whisperParamsMock);
-
-        expect(result.length).toBeGreaterThan(0);
+        let { res } = await whisperAsync(whisperParamsMock, ({ index }) => {
+          console.log(index)
+        });
+
+        expect(res.length).toBeGreaterThan(0);
     }, 10000);
 });
 
diff --git a/examples/addon.node/addon.cpp b/examples/addon.node/addon.cpp
@@ -257,31 +257,56 @@ int run(whisper_params &params, std::vector<std::vector<std::string>> &result) {
     return 0;
 }
 
-class Worker : public Napi::AsyncWorker {
+class Worker : public Napi::AsyncProgressWorker<size_t> {
  public:
   Worker(Napi::Function& callback, whisper_params params)
-      : Napi::AsyncWorker(callback), params(params) {}
+      : Napi::AsyncProgressWorker<size_t>(callback), params(params) {}
 
-  void Execute() override {
-    run(params, result);
+  void Execute(const ExecutionProgress& progress) override {
+    for (size_t i = 0; i < (int) params.fname_inp.size(); ++i) {
+
+      std::vector<std::vector<std::string>> current_result;
+      whisper_params single_file_params = params;
+      single_file_params.fname_inp.clear();
+      single_file_params.fname_inp.push_back(params.fname_inp[i]);
+
+      run(single_file_params, current_result);
+
+      results.push_back(current_result);
+
+      progress.Send(&i, 1);
+    }
+  }
+
+  void OnProgress(const size_t* data, size_t count) override {
+    Napi::HandleScope scope(Env());
+    Napi::Object _res = Napi::Object::New(Env());
+    _res.Set("res", Env().Null());
+    _res.Set("index", Napi::Number::New(Env(), *data));
+    Callback().Call({Env().Null(), _res});
   }
 
   void OnOK() override {
     Napi::HandleScope scope(Env());
-    Napi::Object res = Napi::Array::New(Env(), result.size());
-    for (uint64_t i = 0; i < result.size(); ++i) {
+    Napi::Object res = Napi::Array::New(Env(), results.size());
+    for (uint64_t i = 0; i < results.size(); ++i) {
       Napi::Object tmp = Napi::Array::New(Env(), 3);
       for (uint64_t j = 0; j < 3; ++j) {
-        tmp[j] = Napi::String::New(Env(), result[i][j]);
+        for (uint64_t k = 0; k < 1; ++k) {
+          tmp[j] = Napi::String::New(Env(), results[i][k][j]);
+        }
       }
       res[i] = tmp;
     }
-    Callback().Call({Env().Null(), res});
+    Napi::Object _res = Napi::Object::New(Env());
+    _res.Set("res", res);
+    _res.Set("index", Env().Null());
+    Callback().Call({Env().Null(), _res});
   }
 
  private:
   whisper_params params;
-  std::vector<std::vector<std::string>> result;
+  std::vector<std::vector<std::vector<std::string>>> results;
 };
 
 
@@ -296,12 +321,16 @@ Napi::Value whisper(const Napi::CallbackInfo& info) {
   Napi::Object whisper_params = info[0].As<Napi::Object>();
   std::string language = whisper_params.Get("language").As<Napi::String>();
   std::string model = whisper_params.Get("model").As<Napi::String>();
-  std::string input = whisper_params.Get("fname_inp").As<Napi::String>();
+  Napi::Array inputArray = whisper_params.Get("fname_inp").As<Napi::Array>();
+  std::vector<std::string> input;
+  for (uint32_t i = 0; i < inputArray.Length(); i++) {
+      input.push_back(inputArray.Get(i).As<Napi::String>());
+  }
   bool use_gpu = whisper_params.Get("use_gpu").As<Napi::Boolean>();
 
   params.language = language;
   params.model = model;
-  params.fname_inp.emplace_back(input);
+  params.fname_inp = input;
   params.use_gpu = use_gpu;
 
   Napi::Function callback = info[1].As<Napi::Function>();

diff --git a/examples/addon.node/index.js b/examples/addon.node/index.js
@@ -10,7 +10,7 @@ const whisperAsync = promisify(whisper);
 const whisperParams = {
   language: "en",
   model: path.join(__dirname, "../../models/ggml-base.en.bin"),
-  fname_inp: "../../samples/jfk.wav",
+  fname_inp: ["../../samples/jfk.wav"],
   use_gpu: true,
 };
 
@@ -32,6 +32,6 @@ for (const key in params) {
 
 console.log("whisperParams =", whisperParams);
 
-whisperAsync(whisperParams).then((result) => {
-  console.log(`Result from whisper: ${result}`);
+whisperAsync(whisperParams).then(({res}) => {
+  console.log(`Result from whisper: ${res}`);
 });