cui-llama.rn 1.3.6 → 1.4.0

package/android/src/main/jni.cpp

@@ -14,13 +14,14 @@
 #include "llama-context.h"
 #include "gguf.h"
 #include "rn-llama.hpp"
+#include "jni-utils.h"
 
 #define UNUSED(x) (void)(x)
 #define TAG "RNLLAMA_ANDROID_JNI"
 
 #define LOGI(...) __android_log_print(ANDROID_LOG_INFO,     TAG, __VA_ARGS__)
 #define LOGW(...) __android_log_print(ANDROID_LOG_WARN,     TAG, __VA_ARGS__)
-
+#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR,    TAG, __VA_ARGS__)
 static inline int min(int a, int b) {
     return (a < b) ? a : b;
 }
@@ -129,7 +130,7 @@ static inline void pushString(JNIEnv *env, jobject arr, const char *value) {
 // Method to push WritableMap into WritableArray
 static inline void pushMap(JNIEnv *env, jobject arr, jobject value) {
     jclass mapClass = env->FindClass("com/facebook/react/bridge/WritableArray");
-    jmethodID pushMapMethod = env->GetMethodID(mapClass, "pushMap", "(Lcom/facebook/react/bridge/WritableMap;)V");
+    jmethodID pushMapMethod = env->GetMethodID(mapClass, "pushMap", "(Lcom/facebook/react/bridge/ReadableMap;)V");
 
     env->CallVoidMethod(arr, pushMapMethod, value);
 }
@@ -199,7 +200,7 @@ Java_com_rnllama_LlamaContext_modelInfo(
                 continue;
             }
 
-            const std::string value = rnllama::lm_gguf_kv_to_str(ctx, i);
+            const std::string value = lm_gguf_kv_to_str(ctx, i);
             putString(env, info, key, value.c_str());
         }
     }
@@ -234,16 +235,18 @@ Java_com_rnllama_LlamaContext_initContext(
     jint embd_normalize,
     jint n_ctx,
     jint n_batch,
+    jint n_ubatch,
     jint n_threads,
     jint n_gpu_layers, // TODO: Support this
     jboolean flash_attn,
-    jint cache_type_k,
-    jint cache_type_v,
+    jstring cache_type_k,
+    jstring cache_type_v,
     jboolean use_mlock,
     jboolean use_mmap,
     jboolean vocab_only,
     jstring lora_str,
     jfloat lora_scaled,
+    jobject lora_list,
     jfloat rope_freq_base,
     jfloat rope_freq_scale,
     jint pooling_type,
@@ -263,6 +266,7 @@ Java_com_rnllama_LlamaContext_initContext(
     
     defaultParams.n_ctx = n_ctx;
     defaultParams.n_batch = n_batch;
+    defaultParams.n_ubatch = n_ubatch;
 
     if (pooling_type != -1) {
         defaultParams.pooling_type = static_cast<enum llama_pooling_type>(pooling_type);
@@ -285,19 +289,14 @@ Java_com_rnllama_LlamaContext_initContext(
     // defaultParams.n_gpu_layers = n_gpu_layers;
     defaultParams.flash_attn = flash_attn;
 
-    // const char *cache_type_k_chars = env->GetStringUTFChars(cache_type_k, nullptr);
-    // const char *cache_type_v_chars = env->GetStringUTFChars(cache_type_v, nullptr);
-    defaultParams.cache_type_k = (lm_ggml_type) cache_type_k;
-    defaultParams.cache_type_v = (lm_ggml_type) cache_type_v;
+    const char *cache_type_k_chars = env->GetStringUTFChars(cache_type_k, nullptr);
+    const char *cache_type_v_chars = env->GetStringUTFChars(cache_type_v, nullptr);
+    defaultParams.cache_type_k = rnllama::kv_cache_type_from_str(cache_type_k_chars);
+    defaultParams.cache_type_v = rnllama::kv_cache_type_from_str(cache_type_v_chars);
 
     defaultParams.use_mlock = use_mlock;
     defaultParams.use_mmap = use_mmap;
 
-    const char *lora_chars = env->GetStringUTFChars(lora_str, nullptr);
-    if (lora_chars != nullptr && lora_chars[0] != '\0') {
-        defaultParams.lora_adapters.push_back({lora_chars, lora_scaled});
-    }
-
     defaultParams.rope_freq_base = rope_freq_base;
     defaultParams.rope_freq_scale = rope_freq_scale;
 
@@ -331,23 +330,55 @@ Java_com_rnllama_LlamaContext_initContext(
     bool is_model_loaded = llama->loadModel(defaultParams);
 
     env->ReleaseStringUTFChars(model_path_str, model_path_chars);
-    env->ReleaseStringUTFChars(lora_str, lora_chars);
-    // env->ReleaseStringUTFChars(cache_type_k, cache_type_k_chars);
-    // env->ReleaseStringUTFChars(cache_type_v, cache_type_v_chars);
+    env->ReleaseStringUTFChars(cache_type_k, cache_type_k_chars);
+    env->ReleaseStringUTFChars(cache_type_v, cache_type_v_chars);
 
     LOGI("[RNLlama] is_model_loaded %s", (is_model_loaded ? "true" : "false"));
     if (is_model_loaded) {
-      if (embedding && llama_model_has_encoder(llama->model.get()) && llama_model_has_decoder(llama->model.get())) {
-        LOGI("[RNLlama] computing embeddings in encoder-decoder models is not supported");
-        llama_free(llama->ctx.get());
-        return -1;
-      }
-      context_map[(long) llama->ctx.get()] = llama;
+        if (embedding && llama_model_has_encoder(llama->model) && llama_model_has_decoder(llama->model)) {
+            LOGI("[RNLlama] computing embeddings in encoder-decoder models is not supported");
+            llama_free(llama->ctx);
+            return -1;
+        }
+        context_map[(long) llama->ctx] = llama;
     } else {
-      llama_free(llama->ctx.get());
+        llama_free(llama->ctx);
+    }
+
+    std::vector<common_lora_adapter_info> lora;
+    const char *lora_chars = env->GetStringUTFChars(lora_str, nullptr);
+    if (lora_chars != nullptr && lora_chars[0] != '\0') {
+        common_lora_adapter_info la;
+        la.path = lora_chars;
+        la.scale = lora_scaled;
+        lora.push_back(la);
     }
 
-    return reinterpret_cast<jlong>(llama->ctx.get());
+    if (lora_list != nullptr) {
+        // lora_adapters: ReadableArray<ReadableMap>
+        int lora_list_size = readablearray::size(env, lora_list);
+        for (int i = 0; i < lora_list_size; i++) {
+            jobject lora_adapter = readablearray::getMap(env, lora_list, i);
+            jstring path = readablemap::getString(env, lora_adapter, "path", nullptr);
+            if (path != nullptr) {
+                const char *path_chars = env->GetStringUTFChars(path, nullptr);
+                common_lora_adapter_info la;
+                la.path = path_chars;
+                la.scale = readablemap::getFloat(env, lora_adapter, "scaled", 1.0f);
+                lora.push_back(la);
+                env->ReleaseStringUTFChars(path, path_chars);
+            }
+        }
+    }
+    env->ReleaseStringUTFChars(lora_str, lora_chars);
+    int result = llama->applyLoraAdapters(lora);
+    if (result != 0) {
+      LOGI("[RNLlama] Failed to apply lora adapters");
+      llama_free(llama->ctx);
+      return -1;
+    }
+
+    return reinterpret_cast<jlong>(llama->ctx);
 }
 
 
@@ -373,13 +404,13 @@ Java_com_rnllama_LlamaContext_loadModelDetails(
     UNUSED(thiz);
     auto llama = context_map[(long) context_ptr];
 
-    int count = llama_model_meta_count(llama->model.get());
+    int count = llama_model_meta_count(llama->model);
     auto meta = createWriteableMap(env);
     for (int i = 0; i < count; i++) {
         char key[256];
-        llama_model_meta_key_by_index(llama->model.get(), i, key, sizeof(key));
+        llama_model_meta_key_by_index(llama->model, i, key, sizeof(key));
         char val[2048];
-        llama_model_meta_val_str_by_index(llama->model.get(), i, val, sizeof(val));
+        llama_model_meta_val_str_by_index(llama->model, i, val, sizeof(val));
 
         putString(env, meta, key, val);
     }
@@ -387,10 +418,10 @@ Java_com_rnllama_LlamaContext_loadModelDetails(
     auto result = createWriteableMap(env);
 
     char desc[1024];
-    llama_model_desc(llama->model.get(), desc, sizeof(desc));
+    llama_model_desc(llama->model, desc, sizeof(desc));
     putString(env, result, "desc", desc);
-    putDouble(env, result, "size", llama_model_size(llama->model.get()));
-    putDouble(env, result, "nParams", llama_model_n_params(llama->model.get()));
+    putDouble(env, result, "size", llama_model_size(llama->model));
+    putDouble(env, result, "nParams", llama_model_n_params(llama->model));
     putBoolean(env, result, "isChatTemplateSupported", llama->validateModelChatTemplate());
     putMap(env, result, "metadata", meta);
 
@@ -432,7 +463,7 @@ Java_com_rnllama_LlamaContext_getFormattedChat(
     }
 
     const char *tmpl_chars = env->GetStringUTFChars(chat_template, nullptr);
-    std::string formatted_chat = common_chat_apply_template(llama->model.get(), tmpl_chars, chat, true);
+    std::string formatted_chat = common_chat_apply_template(llama->model, tmpl_chars, chat, true);
 
     return env->NewStringUTF(formatted_chat.c_str());
 }
@@ -451,7 +482,7 @@ Java_com_rnllama_LlamaContext_loadSession(
     auto result = createWriteableMap(env);
     size_t n_token_count_out = 0;
     llama->embd.resize(llama->params.n_ctx);
-    if (!llama_state_load_file(llama->ctx.get(), path_chars, llama->embd.data(), llama->embd.capacity(), &n_token_count_out)) {
+    if (!llama_state_load_file(llama->ctx, path_chars, llama->embd.data(), llama->embd.capacity(), &n_token_count_out)) {
       env->ReleaseStringUTFChars(path, path_chars);
 
       putString(env, result, "error", "Failed to load session");
@@ -460,7 +491,7 @@ Java_com_rnllama_LlamaContext_loadSession(
     llama->embd.resize(n_token_count_out);
     env->ReleaseStringUTFChars(path, path_chars);
 
-    const std::string text = rnllama::tokens_to_str(llama->ctx.get(), llama->embd.cbegin(), llama->embd.cend());
+    const std::string text = rnllama::tokens_to_str(llama->ctx, llama->embd.cbegin(), llama->embd.cend());
     putInt(env, result, "tokens_loaded", n_token_count_out);
     putString(env, result, "prompt", text.c_str());
     return reinterpret_cast<jobject>(result);
@@ -482,7 +513,7 @@ Java_com_rnllama_LlamaContext_saveSession(
     std::vector<llama_token> session_tokens = llama->embd;
     int default_size = session_tokens.size();
     int save_size = size > 0 && size <= default_size ? size : default_size;
-    if (!llama_state_save_file(llama->ctx.get(), path_chars, session_tokens.data(), save_size)) {
+    if (!llama_state_save_file(llama->ctx, path_chars, session_tokens.data(), save_size)) {
       env->ReleaseStringUTFChars(path, path_chars);
       return -1;
     }
@@ -500,13 +531,13 @@ static inline jobject tokenProbsToMap(
     for (const auto &prob : probs) {
         auto probsForToken = createWritableArray(env);
         for (const auto &p : prob.probs) {
-            std::string tokStr = rnllama::tokens_to_output_formatted_string(llama->ctx.get(), p.tok);
+            std::string tokStr = rnllama::tokens_to_output_formatted_string(llama->ctx, p.tok);
             auto probResult = createWriteableMap(env);
             putString(env, probResult, "tok_str", tokStr.c_str());
             putDouble(env, probResult, "prob", p.prob);
             pushMap(env, probsForToken, probResult);
         }
-        std::string tokStr = rnllama::tokens_to_output_formatted_string(llama->ctx.get(), prob.tok);
+        std::string tokStr = rnllama::tokens_to_output_formatted_string(llama->ctx, prob.tok);
         auto tokenResult = createWriteableMap(env);
         putString(env, tokenResult, "content", tokStr.c_str());
         putArray(env, tokenResult, "probs", probsForToken);
@@ -533,7 +564,6 @@ Java_com_rnllama_LlamaContext_doCompletion(
     jfloat mirostat,
     jfloat mirostat_tau,
     jfloat mirostat_eta,
-    jboolean penalize_nl,
     jint top_k,
     jfloat top_p,
     jfloat min_p,
@@ -546,7 +576,7 @@ Java_com_rnllama_LlamaContext_doCompletion(
     jobjectArray logit_bias,
     jfloat   dry_multiplier,
     jfloat   dry_base,
-    jint dry_allowed_length,    
+    jint dry_allowed_length,
     jint dry_penalty_last_n,
     jobjectArray dry_sequence_breakers,
     jobject partial_completion_callback
@@ -556,7 +586,7 @@ Java_com_rnllama_LlamaContext_doCompletion(
 
     llama->rewind();
 
-    //llama_reset_timings(llama->ctx.get());
+    //llama_reset_timings(llama->ctx);
 
     llama->params.prompt = env->GetStringUTFChars(prompt, nullptr);
     llama->params.sampling.seed = (seed == -1) ? time(NULL) : seed;
@@ -578,7 +608,6 @@ Java_com_rnllama_LlamaContext_doCompletion(
     sparams.mirostat = mirostat;
     sparams.mirostat_tau = mirostat_tau;
     sparams.mirostat_eta = mirostat_eta;
-    // sparams.penalize_nl = penalize_nl;
     sparams.top_k = top_k;
     sparams.top_p = top_p;
     sparams.min_p = min_p;
@@ -594,7 +623,7 @@ Java_com_rnllama_LlamaContext_doCompletion(
 
     sparams.logit_bias.clear();
     if (ignore_eos) {
-        sparams.logit_bias[llama_token_eos(llama->model.get())].bias = -INFINITY;
+        sparams.logit_bias[llama_token_eos(llama->model)].bias = -INFINITY;
     }
 
     // dry break seq
@@ -613,7 +642,7 @@ Java_com_rnllama_LlamaContext_doCompletion(
     sparams.dry_sequence_breakers = dry_sequence_breakers_vector;
 
     // logit bias
-    const int n_vocab = llama_n_vocab(llama_get_model(llama->ctx.get()));
+    const int n_vocab = llama_n_vocab(llama_get_model(llama->ctx));
     jsize logit_bias_len = env->GetArrayLength(logit_bias);
 
     for (jsize i = 0; i < logit_bias_len; i++) {
@@ -660,7 +689,7 @@ Java_com_rnllama_LlamaContext_doCompletion(
         if (token_with_probs.tok == -1 || llama->incomplete) {
             continue;
         }
-        const std::string token_text = common_token_to_piece(llama->ctx.get(), token_with_probs.tok);
+        const std::string token_text = common_token_to_piece(llama->ctx, token_with_probs.tok);
 
         size_t pos = std::min(sent_count, llama->generated_text.size());
 
@@ -695,7 +724,7 @@ Java_com_rnllama_LlamaContext_doCompletion(
             putString(env, tokenResult, "token", to_send.c_str());
 
             if (llama->params.sampling.n_probs > 0) {
-              const std::vector<llama_token> to_send_toks = common_tokenize(llama->ctx.get(), to_send, false);
+              const std::vector<llama_token> to_send_toks = common_tokenize(llama->ctx, to_send, false);
               size_t probs_pos = std::min(sent_token_probs_index, llama->generated_token_probs.size());
               size_t probs_stop_pos = std::min(sent_token_probs_index + to_send_toks.size(), llama->generated_token_probs.size());
               if (probs_pos < probs_stop_pos) {
@@ -712,7 +741,7 @@ Java_com_rnllama_LlamaContext_doCompletion(
         }
     }
 
-    llama_perf_context_print(llama->ctx.get());
+    llama_perf_context_print(llama->ctx);
     llama->is_predicting = false;
 
     auto result = createWriteableMap(env);
@@ -727,7 +756,7 @@ Java_com_rnllama_LlamaContext_doCompletion(
     putString(env, result, "stopping_word", llama->stopping_word.c_str());
     putInt(env, result, "tokens_cached", llama->n_past);
 
-    const auto timings_token = llama_perf_context(llama -> ctx.get());
+    const auto timings_token = llama_perf_context(llama -> ctx);
     
     auto timingsResult = createWriteableMap(env);
     putInt(env, timingsResult, "prompt_n", timings_token.n_p_eval);
@@ -771,7 +800,7 @@ Java_com_rnllama_LlamaContext_tokenize(
     const char *text_chars = env->GetStringUTFChars(text, nullptr);
 
     const std::vector<llama_token> toks = common_tokenize(
-        llama->ctx.get(),
+        llama->ctx,
         text_chars,
         false
     );
@@ -798,7 +827,7 @@ Java_com_rnllama_LlamaContext_detokenize(
         toks.push_back(tokens_ptr[i]);
     }
 
-    auto text = rnllama::tokens_to_str(llama->ctx.get(), toks.cbegin(), toks.cend());
+    auto text = rnllama::tokens_to_str(llama->ctx, toks.cbegin(), toks.cend());
 
     env->ReleaseIntArrayElements(tokens, tokens_ptr, 0);
 
@@ -835,7 +864,7 @@ Java_com_rnllama_LlamaContext_embedding(
 
     llama->rewind();
 
-    llama_perf_context_reset(llama->ctx.get());
+    llama_perf_context_reset(llama->ctx);
 
     llama->params.prompt = text_chars;
 
@@ -861,7 +890,7 @@ Java_com_rnllama_LlamaContext_embedding(
 
     auto promptTokens = createWritableArray(env);
     for (const auto &tok : llama->embd) {
-      pushString(env, promptTokens, common_token_to_piece(llama->ctx.get(), tok).c_str());
+      pushString(env, promptTokens, common_token_to_piece(llama->ctx, tok).c_str());
     }
     putArray(env, result, "prompt_tokens", promptTokens);
 
@@ -885,23 +914,64 @@ Java_com_rnllama_LlamaContext_bench(
     return env->NewStringUTF(result.c_str());
 }
 
+JNIEXPORT jint JNICALL
+Java_com_rnllama_LlamaContext_applyLoraAdapters(
+    JNIEnv *env, jobject thiz, jlong context_ptr, jobjectArray loraAdapters) {
+    UNUSED(thiz);
+    auto llama = context_map[(long) context_ptr];
+
+    // lora_adapters: ReadableArray<ReadableMap>
+    std::vector<common_lora_adapter_info> lora_adapters;
+    int lora_adapters_size = readablearray::size(env, loraAdapters);
+    for (int i = 0; i < lora_adapters_size; i++) {
+        jobject lora_adapter = readablearray::getMap(env, loraAdapters, i);
+        jstring path = readablemap::getString(env, lora_adapter, "path", nullptr);
+        if (path != nullptr) {
+          const char *path_chars = env->GetStringUTFChars(path, nullptr);
+          env->ReleaseStringUTFChars(path, path_chars);
+          float scaled = readablemap::getFloat(env, lora_adapter, "scaled", 1.0f);
+          common_lora_adapter_info la;
+          la.path = path_chars;
+          la.scale = scaled;
+          lora_adapters.push_back(la);
+        }
+    }
+    return llama->applyLoraAdapters(lora_adapters);
+}
+
+JNIEXPORT void JNICALL
+Java_com_rnllama_LlamaContext_removeLoraAdapters(
+    JNIEnv *env, jobject thiz, jlong context_ptr) {
+    UNUSED(env);
+    UNUSED(thiz);
+    auto llama = context_map[(long) context_ptr];
+    llama->removeLoraAdapters();
+}
+
+JNIEXPORT jobject JNICALL
+Java_com_rnllama_LlamaContext_getLoadedLoraAdapters(
+    JNIEnv *env, jobject thiz, jlong context_ptr) {
+    UNUSED(thiz);
+    auto llama = context_map[(long) context_ptr];
+    auto loaded_lora_adapters = llama->getLoadedLoraAdapters();
+    auto result = createWritableArray(env);
+    for (common_lora_adapter_info &la : loaded_lora_adapters) {
+        auto map = createWriteableMap(env);
+        putString(env, map, "path", la.path.c_str());
+        putDouble(env, map, "scaled", la.scale);
+        pushMap(env, result, map);
+    }
+    return result;
+}
+
 JNIEXPORT void JNICALL
 Java_com_rnllama_LlamaContext_freeContext(
         JNIEnv *env, jobject thiz, jlong context_ptr) {
     UNUSED(env);
     UNUSED(thiz);
     auto llama = context_map[(long) context_ptr];
-    if (llama->model.get()) {
-        llama_model_free(llama->model.get());
-    }
-    if (llama->ctx.get()) {
-        llama_free(llama->ctx.get());
-    }
-    /*if (llama->ctx.get()-> != nullptr)
-    {
-        common_sampler_free(llama->ctx.get() -> _sampling);
-    }*/
-    context_map.erase((long) llama->ctx.get());
+    context_map.erase((long) llama->ctx);
+    delete llama;
 }
 
 JNIEXPORT void JNICALL

package/android/src/newarch/java/com/rnllama/RNLlamaModule.java

@@ -103,6 +103,21 @@ public class RNLlamaModule extends NativeRNLlamaSpec {
     rnllama.bench(id, pp, tg, pl, nr, promise);
   }
 
+  @ReactMethod
+  public void applyLoraAdapters(double id, final ReadableArray loraAdapters, final Promise promise) {
+    rnllama.applyLoraAdapters(id, loraAdapters, promise);
+  }
+
+  @ReactMethod
+  public void removeLoraAdapters(double id, final Promise promise) {
+    rnllama.removeLoraAdapters(id, promise);
+  }
+
+  @ReactMethod
+  public void getLoadedLoraAdapters(double id, final Promise promise) {
+    rnllama.getLoadedLoraAdapters(id, promise);
+  }
+
   @ReactMethod
   public void releaseContext(double id, Promise promise) {
     rnllama.releaseContext(id, promise);

package/android/src/oldarch/java/com/rnllama/RNLlamaModule.java

@@ -104,6 +104,21 @@ public class RNLlamaModule extends ReactContextBaseJavaModule {
     rnllama.bench(id, pp, tg, pl, nr, promise);
   }
 
+  @ReactMethod
+  public void applyLoraAdapters(double id, final ReadableArray loraAdapters, final Promise promise) {
+    rnllama.applyLoraAdapters(id, loraAdapters, promise);
+  }
+
+  @ReactMethod
+  public void removeLoraAdapters(double id, final Promise promise) {
+    rnllama.removeLoraAdapters(id, promise);
+  }
+
+  @ReactMethod
+  public void getLoadedLoraAdapters(double id, final Promise promise) {
+    rnllama.getLoadedLoraAdapters(id, promise);
+  }
+
   @ReactMethod
   public void releaseContext(double id, Promise promise) {
     rnllama.releaseContext(id, promise);