Load with defauls params

Signed-off-by: Ramya <[email protected]>
rum1887 · Apr 8, 2024 · 841daf7 · 841daf7
1 parent c30c2e5
commit 841daf7
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 3 deletions.
diff --git a/plugins/wasi_nn/wasinnenv.cpp b/plugins/wasi_nn/wasinnenv.cpp
@@ -28,7 +28,7 @@ std::map<std::string_view, Backend> BackendMap = {
     {"tensorflowlite"sv, Backend::TensorflowLite},
     {"autodetect"sv, Backend::Autodetect},
     {"ggml"sv, Backend::GGML},
-    {"whisper"sv, Backend::GGML}};
+    {"whisper"sv, Backend::WHISPER}};
 
 std::map<std::string_view, Device> DeviceMap = {{"cpu"sv, Device::CPU},
                                                 {"gpu"sv, Device::GPU},

diff --git a/plugins/wasi_nn/whispercpp.cpp b/plugins/wasi_nn/whispercpp.cpp
@@ -20,7 +20,51 @@ Expect<ErrNo> load([[maybe_unused]] WasiNNEnvironment &Env,
                    [[maybe_unused]] Span<const Span<uint8_t>> Builders,
                    [[maybe_unused]] Device Device,
                    [[maybe_unused]] uint32_t &GraphId) noexcept {
-  // Env.NNGraph.emplace_back(Backend::WHISPER);
+  Env.NNGraph.emplace_back(Backend::WHISPER);
+  auto &GraphRef = Env.NNGraph.back().get<Graph>();
+  truct whisper_context_params ContextDefault =
+      whisper_context_default_params(); // from whisper.cpp
+  GraphRef.EnableLog = false;
+  GraphRef.EnableDebugLog = false;
+  GraphRef.StreamStdout = false;
+  auto Weight = Builders[0];
+  const std::string BinModel(reinterpret_cast<char *>(Weight.data()),
+                             Weight.size());
+
+  // Handle the model path.
+  std::string ModelFilePath;
+  if (BinModel.substr(0, 8) == "preload:") {
+    ModelFilePath = BinModel.substr(8);
+  } else {
+    // Write whisper model to file.
+    ModelFilePath = "models/ggml-base.en.bin"sv;
+    std::ofstream TempFile(ModelFilePath);
+    if (!TempFile) {
+      spdlog::error(
+          "[WASI-NN] Whisper backend: Failed to create the temporary file.");
+      Env.NNGraph.pop_back();
+      return ErrNo::InvalidArgument;
+    }
+    TempFile << BinModel;
+    TempFile.close();
+  }
+  // auto==struct whisper_context definition available in whisper.cpp which is
+  // not included
+  auto wctx = whisper_init_from_file_with_params(GraphRef.ModelFilePath.c_str(),
+                                                 ContextDefault);
+  GraphRef.WhisperModel = wctx.model;
+
+  if (GraphRef.WhisperModel == nullptr) {
+    spdlog::error("[WASI-NN] Whisper backend: Error: unable to init model."sv);
+    Env.NNGraph.pop_back();
+    return ErrNo::InvalidArgument;
+  }
+  if (GraphRef.EnableDebugLog) {
+    spdlog::info(
+        "[WASI-NN][Debug] Whisper backend: Initialize whisper model with given parameters...Done"sv);
+  }
+  // Store the loaded graph.
+  GraphId = Env.NNGraph.size() - 1;
   return ErrNo::Success;
 }
 
@@ -68,7 +112,7 @@ finiSingle([[maybe_unused]] WASINN::WasiNNEnvironment &Env,
 namespace {
 Expect<ErrNo> reportBackendNotSupported() noexcept {
   spdlog::error("[WASI-NN] whisper backend is not built. use "
-                "-WASMEDGE_PLUGIN_WASI_NN_BACKEND=\"whisper\" to build it."sv);
+                "-WASMEDGE_PLUGIN_WASI_NN_BACKEND=\"WHISPER\" to build it."sv);
   return ErrNo::InvalidArgument;
 }
 } // namespace