JNI: apply the model's chat template so instruct models answer properly

app/src/main/cpp/llama_jni.cpp

@@ -8,6 +8,19 @@ struct State {
     llama_context *ctx;
 };
 
+static std::string apply_chat(const llama_model *model, const std::string &user) {
+    const char *tmpl = llama_model_chat_template(model, nullptr);
+    if (!tmpl) return user;  // base model: use the raw prompt
+    llama_chat_message msg{"user", user.c_str()};
+    std::vector<char> buf(user.size() * 2 + 1024);
+    int n = llama_chat_apply_template(tmpl, &msg, 1, true, buf.data(), (int) buf.size());
+    if (n > (int) buf.size()) {
+        buf.resize(n);
+        n = llama_chat_apply_template(tmpl, &msg, 1, true, buf.data(), (int) buf.size());
+    }
+    return n > 0 ? std::string(buf.data(), n) : user;
+}
+
 extern "C" JNIEXPORT jlong JNICALL
 Java_io_github_abrarul_modularai_engine_LlamaEngine_nativeLoad(JNIEnv *env, jobject, jstring path,
                                                                jint nCtx, jint threads) {
@@ -41,8 +54,9 @@ Java_io_github_abrarul_modularai_engine_LlamaEngine_nativeGenerate(JNIEnv *env,
     const llama_vocab *vocab = llama_model_get_vocab(s->model);
 
     const char *pr = env->GetStringUTFChars(prompt, nullptr);
-    std::string text(pr);
+    std::string raw(pr);
     env->ReleaseStringUTFChars(prompt, pr);
+    std::string text = apply_chat(s->model, raw);
 
     llama_memory_clear(llama_get_memory(s->ctx), true);
 
@@ -62,6 +76,13 @@ Java_io_github_abrarul_modularai_engine_LlamaEngine_nativeGenerate(JNIEnv *env,
 
     jclass cls = env->GetObjectClass(sink);
     jmethodID mid = env->GetMethodID(cls, "onToken", "([B)Z");
+    if (!mid) {
+        llama_sampler_free(smpl);
+        return;
+    }
+
+    int budget = (int) llama_n_ctx(s->ctx) - n;
+    if (maxTokens > budget) maxTokens = budget;
 
     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());