353 lines
14 KiB
C++
353 lines
14 KiB
C++
// Summarize-LLM thin JNI layer over llama.cpp (milestone 1: post-meeting minutes).
|
|
//
|
|
// This is the ONLY native code added for the summarize feature. It is compiled
|
|
// together with a pinned llama.cpp checkout into a single prebuilt shared
|
|
// library (libsummarize-llm-jni.so) so that the app's Gradle build keeps using
|
|
// only prebuilt .so files (no externalNativeBuild, no NDK required on the
|
|
// developer machine). See tools/llama-jni/build_arm64.sh.
|
|
//
|
|
// Safety rules (mirrors the project's "missing asset => native exit()" lesson):
|
|
// - Every entry point validates inputs and throws a Java exception instead
|
|
// of crashing; model loading never calls exit().
|
|
// - The caller (Kotlin) MUST check that the GGUF file exists before
|
|
// nativeNewFromFile(); we double-check with access() here as well.
|
|
|
|
#include <jni.h>
|
|
#include <android/log.h>
|
|
|
|
#include <atomic>
|
|
#include <mutex>
|
|
#include <string>
|
|
#include <vector>
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <random>
|
|
#include <unistd.h>
|
|
|
|
#include "llama.h"
|
|
|
|
#define LOG_TAG "SummarizeLlmJni"
|
|
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__)
|
|
#define LOGI(...) __android_log_print(ANDROID_LOG_INFO, LOG_TAG, __VA_ARGS__)
|
|
|
|
namespace {
|
|
|
|
struct Engine {
|
|
llama_model * model = nullptr;
|
|
llama_context * ctx = nullptr;
|
|
const llama_vocab * vocab = nullptr;
|
|
std::mutex mu; // generate() is single-flight
|
|
std::atomic<bool> abort{false};
|
|
};
|
|
|
|
void throwJava(JNIEnv * env, const char * msg) {
|
|
jclass cls = env->FindClass("java/lang/RuntimeException");
|
|
if (cls != nullptr) {
|
|
env->ThrowNew(cls, msg);
|
|
}
|
|
}
|
|
|
|
std::string jstringToUtf8(JNIEnv * env, jstring s) {
|
|
if (s == nullptr) return {};
|
|
const char * c = env->GetStringUTFChars(s, nullptr);
|
|
std::string out = c != nullptr ? c : "";
|
|
if (c != nullptr) env->ReleaseStringUTFChars(s, c);
|
|
return out;
|
|
}
|
|
|
|
ggml_type parseKvType(const std::string & s, ggml_type fallback) {
|
|
if (s == "f16") return GGML_TYPE_F16;
|
|
if (s == "f32") return GGML_TYPE_F32;
|
|
if (s == "q8_0") return GGML_TYPE_Q8_0;
|
|
if (s == "q4_0") return GGML_TYPE_Q4_0;
|
|
if (s == "q4_1") return GGML_TYPE_Q4_1;
|
|
return fallback;
|
|
}
|
|
|
|
// Abort callback for llama_set_abort_callback.
|
|
bool abortCb(void * user) {
|
|
auto * e = static_cast<Engine *>(user);
|
|
return e->abort.load(std::memory_order_relaxed);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
extern "C" {
|
|
|
|
JNIEXPORT jlong JNICALL
|
|
Java_com_k2fsa_sherpa_onnx_LlmEngine_nativeNewFromFile(
|
|
JNIEnv * env, jclass /*clazz*/,
|
|
jstring jPath, jint ctxSize, jint nThreads,
|
|
jstring jTypeK, jstring jTypeV) {
|
|
|
|
const std::string path = jstringToUtf8(env, jPath);
|
|
if (path.empty()) {
|
|
throwJava(env, "nativeNewFromFile: empty model path");
|
|
return 0;
|
|
}
|
|
if (access(path.c_str(), R_OK) != 0) {
|
|
throwJava(env, ("nativeNewFromFile: model file not readable: " + path).c_str());
|
|
return 0;
|
|
}
|
|
|
|
llama_model_params mparams = llama_model_default_params();
|
|
mparams.n_gpu_layers = 0; // CPU only in milestone 1
|
|
mparams.load_mode = LLAMA_LOAD_MODE_MMAP; // page cache can be reclaimed under memory pressure
|
|
|
|
llama_model * model = llama_model_load_from_file(path.c_str(), mparams);
|
|
if (model == nullptr) {
|
|
throwJava(env, ("nativeNewFromFile: llama failed to load GGUF: " + path).c_str());
|
|
return 0;
|
|
}
|
|
|
|
llama_context_params cparams = llama_context_default_params();
|
|
cparams.n_ctx = ctxSize > 0 ? (uint32_t) ctxSize : 8192;
|
|
cparams.n_threads = nThreads > 0 ? nThreads : 4;
|
|
cparams.n_threads_batch = nThreads > 0 ? nThreads : 4;
|
|
cparams.type_k = parseKvType(jstringToUtf8(env, jTypeK), GGML_TYPE_Q8_0);
|
|
cparams.type_v = parseKvType(jstringToUtf8(env, jTypeV), GGML_TYPE_F16);
|
|
|
|
llama_context * ctx = llama_init_from_model(model, cparams);
|
|
if (ctx == nullptr) {
|
|
llama_model_free(model);
|
|
throwJava(env, "nativeNewFromFile: failed to create llama_context");
|
|
return 0;
|
|
}
|
|
|
|
auto * e = new Engine();
|
|
e->model = model;
|
|
e->ctx = ctx;
|
|
e->vocab = llama_model_get_vocab(model);
|
|
llama_set_abort_callback(ctx, abortCb, e);
|
|
|
|
LOGI("model loaded: %s n_ctx=%d type_k=%d type_v=%d", path.c_str(),
|
|
cparams.n_ctx, (int) cparams.type_k, (int) cparams.type_v);
|
|
return (jlong) (uintptr_t) e;
|
|
}
|
|
|
|
JNIEXPORT jstring JNICALL
|
|
Java_com_k2fsa_sherpa_onnx_LlmEngine_nativeGetStringMeta(
|
|
JNIEnv * env, jclass /*clazz*/, jlong ptr, jstring jKey) {
|
|
|
|
auto * e = (Engine *) (uintptr_t) ptr;
|
|
if (e == nullptr || e->model == nullptr) return nullptr;
|
|
const std::string key = jstringToUtf8(env, jKey);
|
|
char buf[8192];
|
|
int32_t n = llama_model_meta_val_str(e->model, key.c_str(), buf, sizeof(buf));
|
|
if (n < 0) return nullptr;
|
|
return env->NewStringUTF(buf);
|
|
}
|
|
|
|
// M2:录音中降线程(给 ASR 让核),空闲/成稿时拉满。与在途 generate 共用
|
|
// 一把 mu:generate 正在跑时本调用会排队到它结束,不会半路改线程数。
|
|
JNIEXPORT void JNICALL
|
|
Java_com_k2fsa_sherpa_onnx_LlmEngine_nativeSetThreads(
|
|
JNIEnv * /*env*/, jclass /*clazz*/, jlong ptr, jint n) {
|
|
auto * e = (Engine *) (uintptr_t) ptr;
|
|
if (e == nullptr || e->ctx == nullptr || n <= 0) return;
|
|
std::lock_guard<std::mutex> lock(e->mu);
|
|
llama_set_n_threads(e->ctx, n, n);
|
|
}
|
|
|
|
// Blocking generation; call from a background thread. Returns generated text
|
|
// (without the prompt). Stops on EOS, maxNewTokens or nativeAbort().
|
|
JNIEXPORT jstring JNICALL
|
|
Java_com_k2fsa_sherpa_onnx_LlmEngine_nativeGenerate(
|
|
JNIEnv * env, jclass /*clazz*/, jlong ptr,
|
|
jstring jPrompt, jint maxNewTokens, jfloat temperature, jint topK, jfloat topP,
|
|
jfloat repeatPenalty, jint seed, jint minNew) {
|
|
|
|
auto * e = (Engine *) (uintptr_t) ptr;
|
|
if (e == nullptr || e->ctx == nullptr) {
|
|
throwJava(env, "nativeGenerate: engine is released");
|
|
return nullptr;
|
|
}
|
|
std::lock_guard<std::mutex> lock(e->mu);
|
|
e->abort.store(false);
|
|
|
|
// KV 清除(2026-09-16 根治):原实现不清 KV,同引擎连续 generate 会复读
|
|
// 上一条输出 -> Kotlin 侧被迫"每次尝试重建引擎"(E2B 每次重载 4~6s,
|
|
// 是翻译单句 30~40s 的主因)。每次生成前清空内存,引擎即可安全复用;
|
|
// 顺带修复纪要流水线多块共享引擎的潜在污染。
|
|
llama_memory_clear(llama_get_memory(e->ctx), true);
|
|
|
|
const std::string prompt = jstringToUtf8(env, jPrompt);
|
|
if (prompt.empty()) {
|
|
throwJava(env, "nativeGenerate: empty prompt");
|
|
return nullptr;
|
|
}
|
|
|
|
// --- tokenize the prompt ---
|
|
// parse_special=true: the Kotlin side renders the Gemma 4 template with
|
|
// literal <|turn>/<turn|> markers, which must become control tokens 105/106.
|
|
std::vector<llama_token> tokens(prompt.size() / 2 + 64);
|
|
int32_t n;
|
|
while ((n = llama_tokenize(e->vocab, prompt.data(), (int32_t) prompt.size(),
|
|
tokens.data(), (int32_t) tokens.size(),
|
|
/*add_special=*/true, /*parse_special=*/true)) < 0) {
|
|
tokens.resize((size_t) (-n) + 32);
|
|
if (tokens.size() > 4 * 1024 * 1024) {
|
|
throwJava(env, "nativeGenerate: prompt too large to tokenize");
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
// --- fit into the context window (left-truncate keeps the newest content,
|
|
// which for our rolling-notes design is the part that must survive) ---
|
|
const uint32_t nCtx = llama_n_ctx(e->ctx);
|
|
const int32_t keep = (int32_t) (nCtx > 64 ? nCtx - 32 : nCtx);
|
|
if (n > keep) {
|
|
LOGE("prompt tokens %d exceed ctx %u, left-truncating", n, nCtx);
|
|
int32_t drop = n - keep;
|
|
// never drop into the middle of multi-byte UTF-8: only relevant on the
|
|
// decode side; tokens are atomic here, so a plain shift is safe.
|
|
tokens.erase(tokens.begin(), tokens.begin() + drop);
|
|
n = keep;
|
|
}
|
|
|
|
// --- prefill in chunks ---
|
|
const int32_t chunk = 512;
|
|
llama_pos pos = 0;
|
|
for (int32_t i = 0; i < n; i += chunk) {
|
|
if (e->abort.load()) return env->NewStringUTF("");
|
|
int32_t len = std::min(chunk, n - i);
|
|
llama_batch b = llama_batch_get_one(tokens.data() + i, len);
|
|
if (llama_decode(e->ctx, b) != 0) {
|
|
throwJava(env, "nativeGenerate: prefill decode failed");
|
|
return nullptr;
|
|
}
|
|
pos += len;
|
|
}
|
|
|
|
// --- decode loop ---
|
|
const llama_token eos = llama_vocab_eos(e->vocab);
|
|
// seed:调用方显式给随机种子(翻译路径每次点击换盐、每次重试错位,让重试
|
|
// 是真正的再抽签);seed<=0 退回旧行为(42+温度错位,摘要路径用)
|
|
std::mt19937 rng(seed > 0 ? (uint32_t) seed
|
|
: (42 + (uint32_t) (temperature * 1000.0f)));
|
|
std::string out;
|
|
out.reserve(1024);
|
|
std::vector<char> piece(64);
|
|
int32_t cur = tokens.back();
|
|
const int maxNew = maxNewTokens > 0 ? maxNewTokens : 512;
|
|
const bool greedy = !(temperature > 0.0f);
|
|
const int k = topK > 0 ? topK : 64; // Gemma 4 recommended: 64
|
|
const float p = topP > 0.0f ? topP : 0.95f; // recommended: ~0.95
|
|
// Repeat penalty (llama.cpp style, unique-token window): 1B 模型在垃圾/
|
|
// 歧义输入上极易进入复读循环。pen<=0 时禁用(摘要路径保持原行为),
|
|
// 翻译路径传 1.1。只惩罚本步已生成过的 token,窗口=最近 64 个生成 token。
|
|
const bool usePen = repeatPenalty > 1.0f;
|
|
std::vector<llama_token> recent;
|
|
std::vector<bool> penalized; // scratch for unique marking
|
|
|
|
for (int i = 0; i < maxNew; ++i) {
|
|
if (e->abort.load()) break;
|
|
llama_batch b = llama_batch_get_one(&cur, 1);
|
|
if (llama_decode(e->ctx, b) != 0) {
|
|
throwJava(env, "nativeGenerate: decode failed");
|
|
return nullptr;
|
|
}
|
|
float * logits = llama_get_logits_ith(e->ctx, 0);
|
|
const int vsize = llama_vocab_n_tokens(e->vocab);
|
|
|
|
// 防空白输出(2026-09-16 真机实锤:极端量化下个别输入首 token 直落
|
|
// EOS,表现为 raw.len=0 的"空输出"):前 minNew 步把 EOS 置 -inf,
|
|
// 强制模型至少吐几个 token;有害输出交给上层质检闸门。
|
|
// minNew<=0 = 关闭(摘要路径用)。
|
|
if (minNew > 0 && i < minNew && eos >= 0 && eos < vsize) {
|
|
logits[eos] = -INFINITY;
|
|
}
|
|
|
|
if (usePen && !recent.empty()) {
|
|
if ((int) recent.size() > 64)
|
|
recent.erase(recent.begin(), recent.begin() + (recent.size() - 64));
|
|
if (penalized.size() < (size_t) vsize)
|
|
penalized.assign((size_t) vsize + 4096, false);
|
|
for (llama_token t : recent) {
|
|
if (t >= 0 && t < vsize && !penalized[t]) {
|
|
penalized[t] = true;
|
|
float & lg = logits[t];
|
|
lg = lg < 0 ? lg * repeatPenalty : lg / repeatPenalty;
|
|
}
|
|
}
|
|
for (llama_token t : recent) if (t >= 0 && t < vsize) penalized[t] = false;
|
|
}
|
|
|
|
// pick next token
|
|
llama_token next;
|
|
if (greedy) {
|
|
next = 0;
|
|
float best = logits[0];
|
|
for (int t = 1; t < vsize; ++t) {
|
|
if (logits[t] > best) { best = logits[t]; next = t; }
|
|
}
|
|
} else {
|
|
// top-k cut (partial sort), then temperature softmax + top-p cut, then sample
|
|
std::vector<std::pair<float, int32_t>> cand;
|
|
cand.reserve(k);
|
|
for (int t = 0; t < vsize; ++t) {
|
|
if (cand.size() < (size_t) k) {
|
|
cand.emplace_back(logits[t], t);
|
|
std::push_heap(cand.begin(), cand.end());
|
|
} else if (logits[t] > cand.front().first) {
|
|
std::pop_heap(cand.begin(), cand.end());
|
|
cand.back() = {logits[t], t};
|
|
std::push_heap(cand.begin(), cand.end());
|
|
}
|
|
}
|
|
float maxv = cand.front().first;
|
|
double sum = 0;
|
|
for (auto & c : cand) { c.first = std::exp((c.first - maxv) / temperature); sum += c.first; }
|
|
std::sort(cand.begin(), cand.end(), std::greater<>());
|
|
std::uniform_real_distribution<double> dist(0.0, sum);
|
|
double target = dist(rng), acc = 0;
|
|
next = cand[0].second;
|
|
for (auto & c : cand) {
|
|
acc += c.first;
|
|
next = c.second;
|
|
if (acc / sum >= p) break;
|
|
}
|
|
}
|
|
|
|
if (eos >= 0 && next == eos) break;
|
|
if (llama_vocab_is_eog(e->vocab, next)) break;
|
|
|
|
int32_t len = llama_token_to_piece(e->vocab, next, piece.data(),
|
|
(int32_t) piece.size(), 0, false);
|
|
if (len > 0) {
|
|
out.append(piece.data(), (size_t) len);
|
|
// Gemma 4 closes its turn with "<turn|>" (id 106). It may not be
|
|
// registered as eog in GGUF metadata, so stop on the string too
|
|
// and strip it from the output.
|
|
if (out.size() >= 7 && out.compare(out.size() - 7, 7, "<turn|>") == 0) {
|
|
out.erase(out.size() - 7);
|
|
break;
|
|
}
|
|
}
|
|
cur = next;
|
|
if (usePen) recent.push_back(next);
|
|
}
|
|
|
|
LOGI("generated %d chars", (int) out.size());
|
|
return env->NewStringUTF(out.c_str());
|
|
}
|
|
|
|
JNIEXPORT void JNICALL
|
|
Java_com_k2fsa_sherpa_onnx_LlmEngine_nativeAbort(JNIEnv * /*env*/, jclass, jlong ptr) {
|
|
auto * e = (Engine *) (uintptr_t) ptr;
|
|
if (e != nullptr) e->abort.store(true);
|
|
}
|
|
|
|
JNIEXPORT void JNICALL
|
|
Java_com_k2fsa_sherpa_onnx_LlmEngine_nativeFree(JNIEnv * /*env*/, jclass, jlong ptr) {
|
|
auto * e = (Engine *) (uintptr_t) ptr;
|
|
if (e == nullptr) return;
|
|
std::lock_guard<std::mutex> lock(e->mu); // wait for an in-flight generate
|
|
if (e->ctx != nullptr) llama_free(e->ctx);
|
|
if (e->model != nullptr) llama_model_free(e->model);
|
|
delete e;
|
|
LOGI("engine released");
|
|
}
|
|
|
|
} // extern "C"
|