@realtimex/node-llama-cpp 0.1.0

package/llama/addon/AddonContext.h

@@ -0,0 +1,85 @@
+#pragma once
+
+#include <mutex>
+
+#include "llama.h"
+#include "napi.h"
+#include "addonGlobals.h"
+#include "AddonSampler.h"
+
+class AddonContext : public Napi::ObjectWrap<AddonContext> {
+    public:
+        AddonModel* model;
+        llama_context_params context_params;
+        llama_context* ctx;
+        llama_batch batch;
+        uint64_t batchMemorySize = 0;
+        bool has_batch = false;
+        int32_t batch_n_tokens = 0;
+        int n_cur = 0;
+
+        uint64_t loadedContextMemorySize = 0;
+        bool contextLoaded = false;
+
+        bool disposed = false;
+
+        AddonContext(const Napi::CallbackInfo& info);
+        ~AddonContext();
+
+        void dispose();
+        void disposeBatch();
+
+        Napi::Value Init(const Napi::CallbackInfo& info);
+        Napi::Value Dispose(const Napi::CallbackInfo& info);
+
+        Napi::Value GetContextSize(const Napi::CallbackInfo& info);
+        Napi::Value InitBatch(const Napi::CallbackInfo& info);
+        Napi::Value DisposeBatch(const Napi::CallbackInfo& info);
+        Napi::Value AddToBatch(const Napi::CallbackInfo& info);
+        Napi::Value DisposeSequence(const Napi::CallbackInfo& info);
+        Napi::Value RemoveTokenCellsFromSequence(const Napi::CallbackInfo& info);
+        Napi::Value ShiftSequenceTokenCells(const Napi::CallbackInfo& info);
+        Napi::Value GetSequenceKvCacheMinPosition(const Napi::CallbackInfo& info);
+        Napi::Value GetSequenceKvCacheMaxPosition(const Napi::CallbackInfo& info);
+        Napi::Value DecodeBatch(const Napi::CallbackInfo& info);
+        Napi::Value SampleToken(const Napi::CallbackInfo& info);
+
+        Napi::Value GetEmbedding(const Napi::CallbackInfo& info);
+        Napi::Value GetStateSize(const Napi::CallbackInfo& info);
+        Napi::Value GetThreads(const Napi::CallbackInfo& info);
+        Napi::Value SetThreads(const Napi::CallbackInfo& info);
+
+        Napi::Value SaveSequenceStateToFile(const Napi::CallbackInfo& info);
+        Napi::Value LoadSequenceStateFromFile(const Napi::CallbackInfo& info);
+
+        Napi::Value PrintTimings(const Napi::CallbackInfo& info);
+        Napi::Value EnsureDraftContextIsCompatibleForSpeculative(const Napi::CallbackInfo& info);
+
+        Napi::Value SetLoras(const Napi::CallbackInfo& info);
+        Napi::Value RestoreCheckpoint(const Napi::CallbackInfo& info);
+
+        static void init(Napi::Object exports);
+};
+
+class AddonContextSequenceCheckpoint : public Napi::ObjectWrap<AddonContextSequenceCheckpoint> {
+    public:
+        std::mutex dataMutex;
+        std::vector<uint8_t> data;
+        llama_seq_id sequenceId = 0;
+        std::size_t minPos = 0;
+        std::size_t maxPos = 0;
+
+        AddonContextSequenceCheckpoint(const Napi::CallbackInfo& info);
+        ~AddonContextSequenceCheckpoint();
+
+        Napi::Value Init(const Napi::CallbackInfo& info);
+        Napi::Value Dispose(const Napi::CallbackInfo& info);
+
+        void dispose();
+
+        Napi::Value GetSize(const Napi::CallbackInfo& info);
+        Napi::Value GetMinPos(const Napi::CallbackInfo& info);
+        Napi::Value GetMaxPos(const Napi::CallbackInfo& info);
+
+        static void init(Napi::Object exports);
+};

package/llama/addon/AddonGrammar.cpp

@@ -0,0 +1,92 @@
+#include "addonGlobals.h"
+#include "AddonGrammar.h"
+
+AddonGrammar::AddonGrammar(const Napi::CallbackInfo& info) : Napi::ObjectWrap<AddonGrammar>(info) {
+    grammarCode = info[0].As<Napi::String>().Utf8Value();
+
+    if (info.Length() > 1 && info[1].IsObject()) {
+        Napi::Object options = info[1].As<Napi::Object>();
+
+        if (options.Has("addonExports")) {
+            addonExportsRef = Napi::Persistent(options.Get("addonExports").As<Napi::Object>());
+            hasAddonExportsRef = true;
+        }
+
+        if (options.Has("rootRuleName")) {
+            rootRuleName = options.Get("rootRuleName").As<Napi::String>().Utf8Value();
+        }
+    }
+
+    auto parsed_grammar = llama_grammar_init_impl(nullptr, grammarCode.c_str(), rootRuleName.c_str(), false, nullptr, 0, nullptr, 0);
+    
+    // will be empty if there are parse errors
+    if (parsed_grammar == nullptr) {
+        Napi::Error::New(info.Env(), "Failed to parse grammar").ThrowAsJavaScriptException();
+        return;
+    }
+
+    llama_grammar_free_impl(parsed_grammar);
+}
+AddonGrammar::~AddonGrammar() {
+    if (hasAddonExportsRef) {
+        addonExportsRef.Unref();
+        hasAddonExportsRef = false;
+    }
+}
+
+Napi::Value AddonGrammar::isTextCompatible(const Napi::CallbackInfo& info) {
+    const std::string testText = info[0].As<Napi::String>().Utf8Value();
+
+    auto parsed_grammar = llama_grammar_init_impl(nullptr, grammarCode.c_str(), rootRuleName.c_str(), false, nullptr, 0, nullptr, 0);
+    
+    // will be empty if there are parse errors
+    if (parsed_grammar == nullptr) {
+        Napi::Error::New(info.Env(), "Failed to parse grammar").ThrowAsJavaScriptException();
+        return Napi::Boolean::New(info.Env(), false);
+    }
+
+    const auto cpts = unicode_cpts_from_utf8(testText);
+    llama_grammar_stacks & stacks_cur = llama_grammar_get_stacks(parsed_grammar);
+
+    for (const auto & cpt : cpts) {
+        try {
+            llama_grammar_accept(parsed_grammar, cpt);
+        } catch (const std::exception & e) {
+            llama_grammar_free_impl(parsed_grammar);
+            return Napi::Boolean::New(info.Env(), false);
+        } catch (...) {
+            llama_grammar_free_impl(parsed_grammar);
+            return Napi::Boolean::New(info.Env(), false);
+        }
+
+        if (stacks_cur.empty()) {
+            // no stacks means that the grammar failed to match at this point
+            llama_grammar_free_impl(parsed_grammar);
+            return Napi::Boolean::New(info.Env(), false);
+        }
+    }
+
+    for (const auto & stack : stacks_cur) {
+        if (stack.empty()) {
+            // an empty stack means that the grammar has been completed
+            llama_grammar_free_impl(parsed_grammar);
+            return Napi::Boolean::New(info.Env(), true);
+        }
+    }
+
+    llama_grammar_free_impl(parsed_grammar);
+    return Napi::Boolean::New(info.Env(), false);
+}
+
+void AddonGrammar::init(Napi::Object exports) {
+    exports.Set(
+        "AddonGrammar",
+        DefineClass(
+            exports.Env(),
+            "AddonGrammar",
+            {
+                InstanceMethod("isTextCompatible", &AddonGrammar::isTextCompatible),
+            }
+        )
+    );
+}

package/llama/addon/AddonGrammar.h

@@ -0,0 +1,22 @@
+#pragma once
+#include "llama.h"
+#include "common/common.h"
+#include "llama-grammar.h"
+#include "src/unicode.h"
+#include "napi.h"
+#include "addonGlobals.h"
+
+class AddonGrammar : public Napi::ObjectWrap<AddonGrammar> {
+    public:
+        std::string grammarCode = "";
+        std::string rootRuleName = "root";
+        Napi::Reference<Napi::Object> addonExportsRef;
+        bool hasAddonExportsRef = false;
+
+        AddonGrammar(const Napi::CallbackInfo& info);
+        ~AddonGrammar();
+
+        Napi::Value isTextCompatible(const Napi::CallbackInfo& info);
+
+        static void init(Napi::Object exports);
+};

package/llama/addon/AddonGrammarEvaluationState.cpp

@@ -0,0 +1,36 @@
+#include <sstream>
+#include "addonGlobals.h"
+#include "common/common.h"
+#include "llama.h"
+#include "AddonGrammarEvaluationState.h"
+#include "AddonGrammar.h"
+
+AddonGrammarEvaluationState::AddonGrammarEvaluationState(const Napi::CallbackInfo& info) : Napi::ObjectWrap<AddonGrammarEvaluationState>(info) {
+    if (info.Length() == 1) {
+        AddonGrammarEvaluationState* existingState = Napi::ObjectWrap<AddonGrammarEvaluationState>::Unwrap(info[0].As<Napi::Object>());
+        model = existingState->model;
+        model->Ref();
+
+        grammarDef = existingState->grammarDef;
+        grammarDef->Ref();
+
+        sampler = llama_sampler_clone(existingState->sampler);
+    } else {
+        model = Napi::ObjectWrap<AddonModel>::Unwrap(info[0].As<Napi::Object>());
+        model->Ref();
+
+        grammarDef = Napi::ObjectWrap<AddonGrammar>::Unwrap(info[1].As<Napi::Object>());
+        grammarDef->Ref();
+
+        sampler = llama_sampler_init_grammar(model->vocab, grammarDef->grammarCode.c_str(), grammarDef->rootRuleName.c_str());
+    }
+}
+AddonGrammarEvaluationState::~AddonGrammarEvaluationState() {
+    llama_sampler_free(sampler);
+    grammarDef->Unref();
+    model->Unref();
+}
+
+void AddonGrammarEvaluationState::init(Napi::Object exports) {
+    exports.Set("AddonGrammarEvaluationState", DefineClass(exports.Env(), "AddonGrammarEvaluationState", {}));
+}

package/llama/addon/AddonGrammarEvaluationState.h

@@ -0,0 +1,17 @@
+#pragma once
+#include "llama.h"
+#include "napi.h"
+#include "addonGlobals.h"
+#include "AddonModel.h"
+
+class AddonGrammarEvaluationState : public Napi::ObjectWrap<AddonGrammarEvaluationState> {
+    public:
+        AddonModel* model;
+        AddonGrammar* grammarDef;
+        llama_sampler * sampler = nullptr;
+
+        AddonGrammarEvaluationState(const Napi::CallbackInfo& info);
+        ~AddonGrammarEvaluationState();
+
+        static void init(Napi::Object exports);
+};