Phase 3: JNI bridge to llama.cpp (load, streaming generate, free)

app/src/main/cpp/llama_jni.cpp

@@ -0,0 +1,94 @@
+#include <jni.h>
+#include <string>
+#include <vector>
+#include "llama.h"
+
+struct State {
+    llama_model *model;
+    llama_context *ctx;
+};
+
+extern "C" JNIEXPORT jlong JNICALL
+Java_io_github_abrarul_modularai_engine_LlamaEngine_nativeLoad(JNIEnv *env, jobject, jstring path,
+                                                               jint nCtx, jint threads) {
+    llama_backend_init();
+    const char *p = env->GetStringUTFChars(path, nullptr);
+    llama_model_params mp = llama_model_default_params();
+    llama_model *m = llama_model_load_from_file(p, mp);
+    env->ReleaseStringUTFChars(path, p);
+    if (!m) return 0;
+
+    llama_context_params cp = llama_context_default_params();
+    cp.n_ctx = nCtx;
+    cp.n_batch = nCtx;
+    cp.n_ubatch = 512;
+    cp.n_threads = threads;
+    cp.n_threads_batch = threads;
+    llama_context *c = llama_init_from_model(m, cp);
+    if (!c) {
+        llama_model_free(m);
+        return 0;
+    }
+    return reinterpret_cast<jlong>(new State{m, c});
+}
+
+extern "C" JNIEXPORT void JNICALL
+Java_io_github_abrarul_modularai_engine_LlamaEngine_nativeGenerate(JNIEnv *env, jobject, jlong h,
+                                                                   jstring prompt, jint maxTokens,
+                                                                   jfloat temp, jobject sink) {
+    auto *s = reinterpret_cast<State *>(h);
+    if (!s) return;
+    const llama_vocab *vocab = llama_model_get_vocab(s->model);
+
+    const char *pr = env->GetStringUTFChars(prompt, nullptr);
+    std::string text(pr);
+    env->ReleaseStringUTFChars(prompt, pr);
+
+    llama_memory_clear(llama_get_memory(s->ctx), true);
+
+    int n = -llama_tokenize(vocab, text.c_str(), (int) text.size(), nullptr, 0, true, true);
+    if (n <= 0 || n >= (int) llama_n_ctx(s->ctx)) return;  // empty or does not fit the context
+    std::vector<llama_token> toks(n);
+    if (llama_tokenize(vocab, text.c_str(), (int) text.size(), toks.data(), n, true, true) < 0) return;
+
+    llama_sampler *smpl = llama_sampler_chain_init(llama_sampler_chain_default_params());
+    if (temp > 0.0f) {
+        llama_sampler_chain_add(smpl, llama_sampler_init_top_p(0.9f, 1));
+        llama_sampler_chain_add(smpl, llama_sampler_init_temp(temp));
+        llama_sampler_chain_add(smpl, llama_sampler_init_dist(LLAMA_DEFAULT_SEED));
+    } else {
+        llama_sampler_chain_add(smpl, llama_sampler_init_greedy());
+    }
+
+    jclass cls = env->GetObjectClass(sink);
+    jmethodID mid = env->GetMethodID(cls, "onToken", "([B)Z");
+
+    llama_token tok;  // must outlive the loop body: the batch points at it
+    llama_batch batch = llama_batch_get_one(toks.data(), (int) toks.size());
+    for (int i = 0; i < maxTokens; i++) {
+        if (llama_decode(s->ctx, batch) != 0) break;
+        tok = llama_sampler_sample(smpl, s->ctx, -1);
+        if (llama_vocab_is_eog(vocab, tok)) break;
+
+        char buf[256];
+        int len = llama_token_to_piece(vocab, tok, buf, sizeof(buf), 0, true);
+        if (len < 0) break;
+        jbyteArray arr = env->NewByteArray(len);
+        env->SetByteArrayRegion(arr, 0, len, reinterpret_cast<const jbyte *>(buf));
+        jboolean cont = env->CallBooleanMethod(sink, mid, arr);
+        env->DeleteLocalRef(arr);
+        if (!cont || env->ExceptionCheck()) break;
+
+        batch = llama_batch_get_one(&tok, 1);
+    }
+    llama_sampler_free(smpl);
+}
+
+extern "C" JNIEXPORT void JNICALL
+Java_io_github_abrarul_modularai_engine_LlamaEngine_nativeFree(JNIEnv *, jobject, jlong h) {
+    auto *s = reinterpret_cast<State *>(h);
+    if (!s) return;
+    llama_free(s->ctx);
+    llama_model_free(s->model);
+    delete s;
+}