@fugood/llama.node 0.0.1-alpha.4 → 0.2.0

package/CMakeLists.txt

@@ -26,7 +26,7 @@ string(REPLACE "i686" "ia32" ARCH ${ARCH})
 string(REPLACE "i386" "ia32" ARCH ${ARCH})
 string(REPLACE "armv7l" "arm" ARCH ${ARCH})
 string(REPLACE "arm" "arm" ARCH ${ARCH})
-string(REPLACE "arm64ex" "arm64" ARCH ${ARCH})
+string(REPLACE "arm64x" "arm64" ARCH ${ARCH})
 string(REPLACE "aarch64" "arm64" ARCH ${ARCH})
 
 if(DEFINED VARIANT)
@@ -58,6 +58,12 @@ include_directories(${CMAKE_JS_INC})
 # flags: -fPIC
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 
+# VULKAN_SDK
+if (VULKAN_SDK)
+  set(ENV{VULKAN_SDK} ${VULKAN_SDK})
+  find_package(Vulkan REQUIRED)
+endif()
+
 set(LLAMA_STATIC ON CACHE BOOL "Build llama as static library")
 add_subdirectory("src/llama.cpp")
 
@@ -71,6 +77,12 @@ file(
     "src/LlamaCompletionWorker.h"
     "src/LlamaContext.cpp"
     "src/LlamaContext.h"
+    "src/TokenizeWorker.cpp"
+    "src/TokenizeWorker.h"
+    "src/DetokenizeWorker.cpp"
+    "src/DetokenizeWorker.h"
+    "src/EmbeddingWorker.cpp"
+    "src/EmbeddingWorker.h"
     "src/LoadSessionWorker.cpp"
     "src/LoadSessionWorker.h"
     "src/SaveSessionWorker.cpp"
@@ -81,31 +93,54 @@ add_library(${PROJECT_NAME} SHARED ${SOURCE_FILES} ${CMAKE_JS_SRC})
 set_target_properties(${PROJECT_NAME} PROPERTIES PREFIX "" SUFFIX ".node")
 target_link_libraries(${PROJECT_NAME} ${CMAKE_JS_LIB} llama ggml common)
 
+add_custom_target(copy_assets ALL DEPENDS ${PROJECT_NAME})
+
+add_custom_command(
+  TARGET copy_assets
+  COMMAND ${CMAKE_COMMAND} -E remove_directory ${PLATFORM_BINARY_DIR}
+  COMMENT "Cleaning bin folder"
+)
+
 if(MSVC AND CMAKE_JS_NODELIB_DEF AND CMAKE_JS_NODELIB_TARGET)
   # Generate node.lib
   execute_process(COMMAND ${CMAKE_AR} /def:${CMAKE_JS_NODELIB_DEF} /out:${CMAKE_JS_NODELIB_TARGET} ${CMAKE_STATIC_LINKER_FLAGS})
   # copy target to bin folder
   get_filename_component(CMAKE_JS_NODELIB_TARGET_NAME ${CMAKE_JS_NODELIB_TARGET} NAME)
-  add_custom_command(TARGET ${PROJECT_NAME} POST_BUILD
+  add_custom_command(TARGET copy_assets
     COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_JS_NODELIB_TARGET} ${PLATFORM_BINARY_DIR}/${CMAKE_JS_NODELIB_TARGET_NAME}
     COMMENT "Copying to bin folder"
   )
 endif()
 
 # copy target to bin folder
-add_custom_command(TARGET ${PROJECT_NAME} POST_BUILD
+add_custom_command(TARGET copy_assets
   COMMAND ${CMAKE_COMMAND} -E copy $<TARGET_FILE:${PROJECT_NAME}> ${PLATFORM_BINARY_DIR}/$<TARGET_FILE_NAME:${PROJECT_NAME}>
   COMMENT "Copying to bin folder"
 )
 
 if (LLAMA_METAL)
   # copy ${CMAKE_BINARY_DIR}/bin/default.metallib
-  add_custom_target(copy_metallib)
   add_custom_command(
-    TARGET copy_metallib
+    TARGET copy_assets
     COMMAND ${CMAKE_COMMAND} -E copy ${CMAKE_BINARY_DIR}/bin/default.metallib ${PLATFORM_BINARY_DIR}/default.metallib
     COMMENT "Copying default.metallib to bin folder"
   )
-  add_dependencies(copy_metallib ggml-metal)
-  add_dependencies(${PROJECT_NAME} copy_metallib)
+  add_dependencies(copy_assets ggml-metal)
+endif()
+
+if (LLAMA_CLBLAST)
+  find_package(CLBlast)
+  if (CLBlast_FOUND)
+    message(STATUS "CLBlast found: ${CLBlast_DIR}")
+    file(
+      GLOB CLBlast_SO_FILES
+      ${CLBlast_DIR}/../../../bin/clblast.dll
+      ${CLBlast_DIR}/../../../lib/libclblast.so
+    )
+    add_custom_command(
+      TARGET copy_assets
+      COMMAND ${CMAKE_COMMAND} -E copy ${CLBlast_SO_FILES} ${PLATFORM_BINARY_DIR}
+      COMMENT "Copying CLBlast SO files to bin folder"
+    )
+  endif()
 endif()

package/README.md

@@ -1,5 +1,9 @@
 # llama.node
 
+[![CI](https://github.com/mybigday/llama.node/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/mybigday/llama.node/actions/workflows/ci.yml)
+[![NPM Version](https://img.shields.io/npm/v/%40fugood%2Fllama.node)](https://www.npmjs.com/package/@fugood/llama.node)
+![NPM Downloads](https://img.shields.io/npm/dw/%40fugood%2Fllama.node)
+
 Node binding of [llama.cpp](https://github.com/ggerganov/llama.cpp).
 
 [llama.cpp](https://github.com/ggerganov/llama.cpp): Inference of [LLaMA](https://arxiv.org/abs/2302.13971) model in pure C/C++
@@ -22,6 +26,7 @@ const context = await loadModel({
   n_ctx: 2048,
   n_gpu_layers: 1, // > 0: enable GPU
   // embedding: true, // use embedding
+  // lib_variant: 'opencl', // Change backend
 })
 
 // Do completion
@@ -40,6 +45,11 @@ const { text, timings } = await context.completion(
 console.log('Result:', text)
 ```
 
+## Lib Variants
+
+- [x] `default`: General usage, not support GPU except macOS (Metal)
+- [x] `vulkan`: Support GPU Vulkan (Windows/Linux), but some scenario might unstable
+
 ## License
 
 MIT

package/lib/binding.js

@@ -50,7 +50,7 @@ const setupEnv = (variant) => {
 };
 const loadModule = (variant) => __awaiter(void 0, void 0, void 0, function* () {
     try {
-        if (variant) {
+        if (variant && variant !== 'default') {
             setupEnv(variant);
             return yield Promise.resolve(`${`../bin/${process.platform}-${variant}/${process.arch}/llama-node.node`}`).then(s => __importStar(require(s)));
         }

package/lib/binding.ts

@@ -23,6 +23,7 @@ export type LlamaCompletionOptions = {
   max_tokens?: number
   seed?: number
   stop?: string[]
+  grammar?: string
 }
 
 export type LlamaCompletionResult = {
@@ -36,11 +37,22 @@ export type LlamaCompletionToken = {
   token: string
 }
 
+export type TokenizeResult = {
+  tokens: Int32Array
+}
+
+export type EmbeddingResult = {
+  embedding: Float32Array
+}
+
 export interface LlamaContext {
   new (options: LlamaModelOptions): LlamaContext
   getSystemInfo(): string
   completion(options: LlamaCompletionOptions, callback?: (token: LlamaCompletionToken) => void): Promise<LlamaCompletionResult>
   stopCompletion(): void
+  tokenize(text: string): Promise<TokenizeResult>
+  detokenize(tokens: number[]): Promise<string>
+  embedding(text: string): Promise<EmbeddingResult>
   saveSession(path: string): Promise<void>
   loadSession(path: string): Promise<void>
   release(): Promise<void>
@@ -50,6 +62,8 @@ export interface Module {
   LlamaContext: LlamaContext
 }
 
+export type LibVariant = 'default' | 'opencl'
+
 const setupEnv = (variant?: string) => {
   const postfix = variant ? `-${variant}` : ''
   const binPath = path.resolve(__dirname, `../bin/${process.platform}${postfix}/${process.arch}/`)
@@ -63,9 +77,9 @@ const setupEnv = (variant?: string) => {
   }
 }
 
-export const loadModule = async (variant?: string): Promise<Module> => {
+export const loadModule = async (variant?: LibVariant): Promise<Module> => {
   try {
-    if (variant) {
+    if (variant && variant !== 'default') {
       setupEnv(variant)
       return await import(`../bin/${process.platform}-${variant}/${process.arch}/llama-node.node`) as Module
     }

package/lib/index.ts

@@ -1,10 +1,10 @@
 import { loadModule, LlamaModelOptions } from './binding'
-import type { Module, LlamaContext } from './binding'
+import type { Module, LlamaContext, LibVariant } from './binding'
 
 export * from './binding'
 
 export interface LlamaModelOptionsExtended extends LlamaModelOptions {
-  lib_variant?: string
+  lib_variant?: LibVariant
 }
 
 const mods: { [key: string]: Module } = {}

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@fugood/llama.node",
   "access": "public",
-  "version": "0.0.1-alpha.4",
+  "version": "0.2.0",
   "description": "Llama.cpp for Node.js",
   "main": "lib/index.js",
   "scripts": {
@@ -39,8 +39,6 @@
   },
   "files": [
     "bin/**/*",
-    "scripts/*.js",
-    "scripts/*.ts",
     "src/**/*.{c,cc,cpp,h,hh,hpp,txt,cmake}",
     "lib/*.js",
     "lib/*.ts",
@@ -62,5 +60,19 @@
   },
   "dependencies": {
     "node-addon-api": "^8.0.0"
+  },
+  "jest": {
+    "testEnvironment": "node",
+    "moduleFileExtensions": [
+      "ts",
+      "tsx",
+      "js",
+      "jsx",
+      "json",
+      "node"
+    ],
+    "testMatch": [
+      "**/*.test.ts"
+    ]
   }
 }

package/src/DetokenizeWorker.cpp

@@ -0,0 +1,22 @@
+#include "DetokenizeWorker.h"
+#include "LlamaContext.h"
+
+DetokenizeWorker::DetokenizeWorker(const Napi::CallbackInfo &info,
+                                   LlamaSessionPtr &sess,
+                                   std::vector<llama_token> &tokens)
+    : AsyncWorker(info.Env()), Deferred(info.Env()), _sess(sess),
+      _tokens(std::move(tokens)) {}
+
+void DetokenizeWorker::Execute() {
+  const auto text = ::llama_detokenize_bpe(_sess->context(), _tokens);
+  _text = std::move(text);
+}
+
+void DetokenizeWorker::OnOK() {
+  Napi::Promise::Deferred::Resolve(
+      Napi::String::New(Napi::AsyncWorker::Env(), _text));
+}
+
+void DetokenizeWorker::OnError(const Napi::Error &err) {
+  Napi::Promise::Deferred::Reject(err.Value());
+}

package/src/DetokenizeWorker.h

@@ -0,0 +1,19 @@
+#include "common.hpp"
+#include <vector>
+
+class DetokenizeWorker : public Napi::AsyncWorker,
+                         public Napi::Promise::Deferred {
+public:
+  DetokenizeWorker(const Napi::CallbackInfo &info, LlamaSessionPtr &sess,
+                   std::vector<llama_token> &tokens);
+
+protected:
+  void Execute();
+  void OnOK();
+  void OnError(const Napi::Error &err);
+
+private:
+  LlamaSessionPtr _sess;
+  std::vector<llama_token> _tokens;
+  std::string _text;
+};

package/src/EmbeddingWorker.cpp

@@ -0,0 +1,46 @@
+#include "EmbeddingWorker.h"
+#include "LlamaContext.h"
+
+EmbeddingWorker::EmbeddingWorker(const Napi::CallbackInfo &info,
+                                 LlamaSessionPtr &sess, std::string text)
+    : AsyncWorker(info.Env()), Deferred(info.Env()), _sess(sess), _text(text) {}
+
+void EmbeddingWorker::Execute() {
+  llama_kv_cache_clear(_sess->context());
+  auto tokens = ::llama_tokenize(_sess->context(), _text, true);
+  // add SEP if not present
+  if (tokens.empty() || tokens.back() != llama_token_sep(_sess->model())) {
+    tokens.push_back(llama_token_sep(_sess->model()));
+  }
+  const int n_embd = llama_n_embd(_sess->model());
+  do {
+    int ret =
+        llama_decode(_sess->context(),
+                     llama_batch_get_one(tokens.data(), tokens.size(), 0, 0));
+    if (ret < 0) {
+      SetError("Failed to inference, code: " + std::to_string(ret));
+      break;
+    }
+    const float *embd = llama_get_embeddings_seq(_sess->context(), 0);
+    if (embd == nullptr) {
+      SetError("Failed to get embeddings");
+      break;
+    }
+    _result.embedding.resize(n_embd);
+    memcpy(_result.embedding.data(), embd, n_embd * sizeof(float));
+  } while (false);
+}
+
+void EmbeddingWorker::OnOK() {
+  auto result = Napi::Object::New(Napi::AsyncWorker::Env());
+  auto embedding = Napi::Float32Array::New(Napi::AsyncWorker::Env(),
+                                           _result.embedding.size());
+  memcpy(embedding.Data(), _result.embedding.data(),
+         _result.embedding.size() * sizeof(float));
+  result.Set("embedding", embedding);
+  Napi::Promise::Deferred::Resolve(result);
+}
+
+void EmbeddingWorker::OnError(const Napi::Error &err) {
+  Napi::Promise::Deferred::Reject(err.Value());
+}

package/src/EmbeddingWorker.h

@@ -0,0 +1,23 @@
+#include "common.hpp"
+#include <vector>
+
+struct EmbeddingResult {
+  std::vector<float> embedding;
+};
+
+class EmbeddingWorker : public Napi::AsyncWorker,
+                        public Napi::Promise::Deferred {
+public:
+  EmbeddingWorker(const Napi::CallbackInfo &info, LlamaSessionPtr &sess,
+                  std::string text);
+
+protected:
+  void Execute();
+  void OnOK();
+  void OnError(const Napi::Error &err);
+
+private:
+  LlamaSessionPtr _sess;
+  std::string _text;
+  EmbeddingResult _result;
+};

package/src/LlamaCompletionWorker.cpp

@@ -58,7 +58,7 @@ void LlamaCompletionWorker::Execute() {
   const auto n_keep = _params.n_keep;
   size_t n_cur = 0;
   size_t n_input = 0;
-  const auto model = llama_get_model(_sess->context());
+  const auto model = _sess->model();
   const bool add_bos = llama_should_add_bos_token(model);
   auto ctx = _sess->context();
 
@@ -110,6 +110,7 @@ void LlamaCompletionWorker::Execute() {
     // sample the next token
     const llama_token new_token_id =
         llama_sampling_sample(sampling.get(), ctx, nullptr);
+    llama_sampling_accept(sampling.get(), ctx, new_token_id, true);
     // prepare the next batch
     embd->emplace_back(new_token_id);
     auto token = llama_token_to_piece(ctx, new_token_id);
@@ -143,6 +144,9 @@ void LlamaCompletionWorker::Execute() {
   }
   const auto t_main_end = ggml_time_us();
   _sess->get_mutex().unlock();
+  if (_onComplete) {
+    _onComplete();
+  }
 }
 
 void LlamaCompletionWorker::OnOK() {

package/src/LlamaCompletionWorker.h

@@ -1,4 +1,5 @@
 #include "common.hpp"
+#include <functional>
 
 struct CompletionResult {
   std::string text = "";
@@ -18,6 +19,8 @@ public:
 
   inline void Stop() { _stop = true; }
 
+  inline void onComplete(std::function<void()> cb) { _onComplete = cb; }
+
 protected:
   void Execute();
   void OnOK();
@@ -30,5 +33,6 @@ private:
   Napi::ThreadSafeFunction _tsfn;
   bool _has_callback = false;
   bool _stop = false;
+  std::function<void()> _onComplete;
   CompletionResult _result;
 };

package/src/LlamaContext.cpp

@@ -1,8 +1,11 @@
 #include "LlamaContext.h"
+#include "DetokenizeWorker.h"
 #include "DisposeWorker.h"
+#include "EmbeddingWorker.h"
 #include "LlamaCompletionWorker.h"
 #include "LoadSessionWorker.h"
 #include "SaveSessionWorker.h"
+#include "TokenizeWorker.h"
 
 void LlamaContext::Init(Napi::Env env, Napi::Object &exports) {
   Napi::Function func = DefineClass(
@@ -16,6 +19,13 @@ void LlamaContext::Init(Napi::Env env, Napi::Object &exports) {
        InstanceMethod<&LlamaContext::StopCompletion>(
            "stopCompletion",
            static_cast<napi_property_attributes>(napi_enumerable)),
+       InstanceMethod<&LlamaContext::Tokenize>(
+           "tokenize", static_cast<napi_property_attributes>(napi_enumerable)),
+       InstanceMethod<&LlamaContext::Detokenize>(
+           "detokenize",
+           static_cast<napi_property_attributes>(napi_enumerable)),
+       InstanceMethod<&LlamaContext::Embedding>(
+           "embedding", static_cast<napi_property_attributes>(napi_enumerable)),
        InstanceMethod<&LlamaContext::SaveSession>(
            "saveSession",
            static_cast<napi_property_attributes>(napi_enumerable)),
@@ -70,7 +80,7 @@ LlamaContext::LlamaContext(const Napi::CallbackInfo &info)
         .ThrowAsJavaScriptException();
   }
 
-  _sess = std::make_shared<LlamaSession>(ctx, params);
+  _sess = std::make_shared<LlamaSession>(model, ctx, params);
   _info = get_system_info(params);
 }
 
@@ -93,6 +103,10 @@ Napi::Value LlamaContext::Completion(const Napi::CallbackInfo &info) {
     Napi::TypeError::New(env, "Context is disposed")
         .ThrowAsJavaScriptException();
   }
+  if (_wip != nullptr) {
+    Napi::TypeError::New(env, "Another completion is in progress")
+        .ThrowAsJavaScriptException();
+  }
   auto options = info[0].As<Napi::Object>();
 
   gpt_params params = _sess->params();
@@ -143,6 +157,7 @@ Napi::Value LlamaContext::Completion(const Napi::CallbackInfo &info) {
       new LlamaCompletionWorker(info, _sess, callback, params, stop_words);
   worker->Queue();
   _wip = worker;
+  worker->onComplete([this]() { _wip = nullptr; });
   return worker->Promise();
 }
 
@@ -153,6 +168,58 @@ void LlamaContext::StopCompletion(const Napi::CallbackInfo &info) {
   }
 }
 
+// tokenize(text: string): Promise<TokenizeResult>
+Napi::Value LlamaContext::Tokenize(const Napi::CallbackInfo &info) {
+  Napi::Env env = info.Env();
+  if (info.Length() < 1 || !info[0].IsString()) {
+    Napi::TypeError::New(env, "String expected").ThrowAsJavaScriptException();
+  }
+  if (_sess == nullptr) {
+    Napi::TypeError::New(env, "Context is disposed")
+        .ThrowAsJavaScriptException();
+  }
+  auto text = info[0].ToString().Utf8Value();
+  auto *worker = new TokenizeWorker(info, _sess, text);
+  worker->Queue();
+  return worker->Promise();
+}
+
+// detokenize(tokens: number[]): Promise<string>
+Napi::Value LlamaContext::Detokenize(const Napi::CallbackInfo &info) {
+  Napi::Env env = info.Env();
+  if (info.Length() < 1 || !info[0].IsArray()) {
+    Napi::TypeError::New(env, "Array expected").ThrowAsJavaScriptException();
+  }
+  if (_sess == nullptr) {
+    Napi::TypeError::New(env, "Context is disposed")
+        .ThrowAsJavaScriptException();
+  }
+  auto tokens = info[0].As<Napi::Array>();
+  std::vector<int32_t> token_ids;
+  for (size_t i = 0; i < tokens.Length(); i++) {
+    token_ids.push_back(tokens.Get(i).ToNumber().Int32Value());
+  }
+  auto *worker = new DetokenizeWorker(info, _sess, token_ids);
+  worker->Queue();
+  return worker->Promise();
+}
+
+// embedding(text: string): Promise<EmbeddingResult>
+Napi::Value LlamaContext::Embedding(const Napi::CallbackInfo &info) {
+  Napi::Env env = info.Env();
+  if (info.Length() < 1 || !info[0].IsString()) {
+    Napi::TypeError::New(env, "String expected").ThrowAsJavaScriptException();
+  }
+  if (_sess == nullptr) {
+    Napi::TypeError::New(env, "Context is disposed")
+        .ThrowAsJavaScriptException();
+  }
+  auto text = info[0].ToString().Utf8Value();
+  auto *worker = new EmbeddingWorker(info, _sess, text);
+  worker->Queue();
+  return worker->Promise();
+}
+
 // saveSession(path: string): Promise<void> throws error
 Napi::Value LlamaContext::SaveSession(const Napi::CallbackInfo &info) {
   Napi::Env env = info.Env();
@@ -163,6 +230,12 @@ Napi::Value LlamaContext::SaveSession(const Napi::CallbackInfo &info) {
     Napi::TypeError::New(env, "Context is disposed")
         .ThrowAsJavaScriptException();
   }
+#ifdef GGML_USE_VULKAN
+  if (_sess->params().n_gpu_layers > 0) {
+    Napi::TypeError::New(env, "Vulkan cannot save session")
+        .ThrowAsJavaScriptException();
+  }
+#endif
   auto *worker = new SaveSessionWorker(info, _sess);
   worker->Queue();
   return worker->Promise();
@@ -178,6 +251,12 @@ Napi::Value LlamaContext::LoadSession(const Napi::CallbackInfo &info) {
     Napi::TypeError::New(env, "Context is disposed")
         .ThrowAsJavaScriptException();
   }
+#ifdef GGML_USE_VULKAN
+  if (_sess->params().n_gpu_layers > 0) {
+    Napi::TypeError::New(env, "Vulkan cannot load session")
+        .ThrowAsJavaScriptException();
+  }
+#endif
   auto *worker = new LoadSessionWorker(info, _sess);
   worker->Queue();
   return worker->Promise();

package/src/LlamaContext.h

@@ -11,6 +11,9 @@ private:
   Napi::Value GetSystemInfo(const Napi::CallbackInfo &info);
   Napi::Value Completion(const Napi::CallbackInfo &info);
   void StopCompletion(const Napi::CallbackInfo &info);
+  Napi::Value Tokenize(const Napi::CallbackInfo &info);
+  Napi::Value Detokenize(const Napi::CallbackInfo &info);
+  Napi::Value Embedding(const Napi::CallbackInfo &info);
   Napi::Value SaveSession(const Napi::CallbackInfo &info);
   Napi::Value LoadSession(const Napi::CallbackInfo &info);
   Napi::Value Release(const Napi::CallbackInfo &info);

package/src/TokenizeWorker.cpp

@@ -0,0 +1,26 @@
+#include "TokenizeWorker.h"
+#include "LlamaContext.h"
+
+TokenizeWorker::TokenizeWorker(const Napi::CallbackInfo &info,
+                               LlamaSessionPtr &sess, std::string text)
+    : AsyncWorker(info.Env()), Deferred(info.Env()), _sess(sess), _text(text) {}
+
+void TokenizeWorker::Execute() {
+  const auto tokens = ::llama_tokenize(_sess->context(), _text, false);
+  _result = {.tokens = std::move(tokens)};
+}
+
+void TokenizeWorker::OnOK() {
+  Napi::HandleScope scope(Napi::AsyncWorker::Env());
+  auto result = Napi::Object::New(Napi::AsyncWorker::Env());
+  auto tokens =
+      Napi::Int32Array::New(Napi::AsyncWorker::Env(), _result.tokens.size());
+  memcpy(tokens.Data(), _result.tokens.data(),
+         _result.tokens.size() * sizeof(llama_token));
+  result.Set("tokens", tokens);
+  Napi::Promise::Deferred::Resolve(result);
+}
+
+void TokenizeWorker::OnError(const Napi::Error &err) {
+  Napi::Promise::Deferred::Reject(err.Value());
+}

package/src/TokenizeWorker.h

@@ -0,0 +1,23 @@
+#include "common.hpp"
+#include <vector>
+
+struct TokenizeResult {
+  std::vector<llama_token> tokens;
+};
+
+class TokenizeWorker : public Napi::AsyncWorker,
+                       public Napi::Promise::Deferred {
+public:
+  TokenizeWorker(const Napi::CallbackInfo &info, LlamaSessionPtr &sess,
+                 std::string text);
+
+protected:
+  void Execute();
+  void OnOK();
+  void OnError(const Napi::Error &err);
+
+private:
+  LlamaSessionPtr _sess;
+  std::string _text;
+  TokenizeResult _result;
+};

package/src/common.hpp

@@ -46,32 +46,37 @@ constexpr T get_option(const Napi::Object &options, const std::string &name,
 
 class LlamaSession {
 public:
-  LlamaSession(llama_context *ctx, gpt_params params)
-      : ctx_(LlamaCppContext(ctx, llama_free)), params_(params) {
+  LlamaSession(llama_model *model, llama_context *ctx, gpt_params params)
+      : model_(LlamaCppModel(model, llama_free_model)),
+        ctx_(LlamaCppContext(ctx, llama_free)), params_(params) {
     tokens_.reserve(params.n_ctx);
   }
 
   ~LlamaSession() { dispose(); }
 
-  llama_context *context() { return ctx_.get(); }
+  inline llama_context *context() { return ctx_.get(); }
 
-  std::vector<llama_token>* tokens_ptr() { return &tokens_; }
+  inline llama_model *model() { return model_.get(); }
 
-  void set_tokens(std::vector<llama_token> tokens) {
+  inline std::vector<llama_token> *tokens_ptr() { return &tokens_; }
+
+  inline void set_tokens(std::vector<llama_token> tokens) {
     tokens_ = std::move(tokens);
   }
 
-  const gpt_params &params() const { return params_; }
+  inline const gpt_params &params() const { return params_; }
 
-  std::mutex &get_mutex() { return mutex; }
+  inline std::mutex &get_mutex() { return mutex; }
 
   void dispose() {
     std::lock_guard<std::mutex> lock(mutex);
     tokens_.clear();
     ctx_.reset();
+    model_.reset();
   }
 
 private:
+  LlamaCppModel model_;
   LlamaCppContext ctx_;
   const gpt_params params_;
   std::vector<llama_token> tokens_{};