342 lines
11 KiB
C++
342 lines
11 KiB
C++
#include "ggml-backend.h"
|
|
#include "llama.h"
|
|
|
|
#include <algorithm>
|
|
#include <cerrno>
|
|
#include <cmath>
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <fstream>
|
|
#include <limits>
|
|
#include <sstream>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
static void die(const char *msg) {
|
|
std::fprintf(stderr, "%s\n", msg);
|
|
std::exit(1);
|
|
}
|
|
|
|
static std::string read_file(const char *path) {
|
|
std::ifstream in(path, std::ios::binary);
|
|
if (!in) {
|
|
std::fprintf(stderr, "open %s: %s\n", path, std::strerror(errno));
|
|
std::exit(1);
|
|
}
|
|
std::ostringstream ss;
|
|
ss << in.rdbuf();
|
|
return ss.str();
|
|
}
|
|
|
|
static void strip_newline(std::string &s) {
|
|
while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) {
|
|
s.pop_back();
|
|
}
|
|
}
|
|
|
|
static std::vector<std::string> split_tab(const std::string &line) {
|
|
std::vector<std::string> out;
|
|
size_t start = 0;
|
|
for (;;) {
|
|
size_t tab = line.find('\t', start);
|
|
if (tab == std::string::npos) {
|
|
out.push_back(line.substr(start));
|
|
return out;
|
|
}
|
|
out.push_back(line.substr(start, tab - start));
|
|
start = tab + 1;
|
|
}
|
|
}
|
|
|
|
static std::vector<llama_token> tokenize(
|
|
const llama_vocab *vocab,
|
|
const std::string &text,
|
|
bool add_special,
|
|
bool parse_special) {
|
|
int n = llama_tokenize(vocab, text.data(), (int32_t)text.size(),
|
|
nullptr, 0, add_special, parse_special);
|
|
if (n < 0) n = -n;
|
|
if (n == 0) return {};
|
|
|
|
std::vector<llama_token> tokens((size_t)n);
|
|
int got = llama_tokenize(vocab, text.data(), (int32_t)text.size(),
|
|
tokens.data(), n, add_special, parse_special);
|
|
if (got < 0) die("llama_tokenize failed");
|
|
tokens.resize((size_t)got);
|
|
return tokens;
|
|
}
|
|
|
|
static std::string render_glm_ds4_prompt(const std::string &prompt) {
|
|
return std::string("[gMASK]<sop><|user|>") + prompt +
|
|
"<|assistant|><think></think>";
|
|
}
|
|
|
|
static std::string render_deepseek_ds4_prompt(const std::string &prompt) {
|
|
return std::string(u8"<\uFF5Cbegin\u2581of\u2581sentence\uFF5C>"
|
|
u8"<\uFF5CUser\uFF5C>") + prompt +
|
|
u8"<\uFF5CAssistant\uFF5C></think>";
|
|
}
|
|
|
|
static std::string render_template_prompt(
|
|
const char *tmpl,
|
|
const std::string &prompt,
|
|
bool *ok) {
|
|
llama_chat_message msg = {"user", prompt.c_str()};
|
|
int n = llama_chat_apply_template(tmpl, &msg, 1, true, nullptr, 0);
|
|
if (n < 0) {
|
|
*ok = false;
|
|
return {};
|
|
}
|
|
std::vector<char> buf((size_t)n + 1);
|
|
int got = llama_chat_apply_template(tmpl, &msg, 1, true,
|
|
buf.data(), (int32_t)buf.size());
|
|
if (got < 0) {
|
|
*ok = false;
|
|
return {};
|
|
}
|
|
*ok = true;
|
|
return std::string(buf.data(), (size_t)got);
|
|
}
|
|
|
|
static bool decode_chunk(
|
|
llama_context *ctx,
|
|
llama_batch &batch,
|
|
const llama_token *tokens,
|
|
int n_tokens,
|
|
int pos,
|
|
bool logits_last) {
|
|
batch.n_tokens = n_tokens;
|
|
for (int i = 0; i < n_tokens; i++) {
|
|
batch.token[i] = tokens[i];
|
|
batch.pos[i] = pos + i;
|
|
batch.n_seq_id[i] = 1;
|
|
batch.seq_id[i][0] = 0;
|
|
batch.logits[i] = (logits_last && i == n_tokens - 1) ? 1 : 0;
|
|
}
|
|
return llama_decode(ctx, batch) == 0;
|
|
}
|
|
|
|
static bool decode_tokens(
|
|
llama_context *ctx,
|
|
llama_batch &batch,
|
|
const std::vector<llama_token> &tokens,
|
|
int start_pos,
|
|
int n_batch,
|
|
bool logits_last) {
|
|
int off = 0;
|
|
while (off < (int)tokens.size()) {
|
|
int n = std::min(n_batch, (int)tokens.size() - off);
|
|
bool want_logits = logits_last && off + n == (int)tokens.size();
|
|
if (!decode_chunk(ctx, batch, tokens.data() + off, n,
|
|
start_pos + off, want_logits)) {
|
|
return false;
|
|
}
|
|
off += n;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
static double token_logprob(
|
|
const float *logits,
|
|
int n_vocab,
|
|
llama_token token,
|
|
llama_token *greedy_out) {
|
|
float max_logit = -std::numeric_limits<float>::infinity();
|
|
llama_token greedy = 0;
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
if (logits[i] > max_logit) {
|
|
max_logit = logits[i];
|
|
greedy = (llama_token)i;
|
|
}
|
|
}
|
|
|
|
double sum = 0.0;
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
sum += std::exp((double)logits[i] - (double)max_logit);
|
|
}
|
|
*greedy_out = greedy;
|
|
return (double)logits[token] - ((double)max_logit + std::log(sum));
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
if (argc != 4 && argc != 5 && argc != 6) {
|
|
std::fprintf(stderr,
|
|
"usage: %s MODEL manifest.tsv OUT.tsv [ctx] "
|
|
"[auto|deepseek-ds4|glm-ds4]\n",
|
|
argv[0]);
|
|
return 2;
|
|
}
|
|
|
|
const char *model_path = argv[1];
|
|
const char *manifest_path = argv[2];
|
|
const char *out_path = argv[3];
|
|
int ctx_size = argc >= 5 ? std::atoi(argv[4]) : 4096;
|
|
if (ctx_size < 1024) ctx_size = 1024;
|
|
const std::string template_mode = argc == 6 ? argv[5] : "auto";
|
|
if (template_mode != "auto" && template_mode != "deepseek-ds4" &&
|
|
template_mode != "glm-ds4") {
|
|
die("template mode must be auto, deepseek-ds4, or glm-ds4");
|
|
}
|
|
|
|
ggml_backend_load_all();
|
|
llama_backend_init();
|
|
|
|
llama_model_params model_params = llama_model_default_params();
|
|
model_params.n_gpu_layers = -1;
|
|
model_params.use_mmap = true;
|
|
|
|
llama_model *model = llama_model_load_from_file(model_path, model_params);
|
|
if (!model) die("failed to open model");
|
|
const llama_vocab *vocab = llama_model_get_vocab(model);
|
|
const int n_vocab = llama_vocab_n_tokens(vocab);
|
|
const char *tmpl = llama_model_chat_template(model, nullptr);
|
|
|
|
llama_context_params ctx_params = llama_context_default_params();
|
|
ctx_params.n_ctx = (uint32_t)ctx_size;
|
|
ctx_params.n_batch = 2048;
|
|
ctx_params.n_ubatch = 512;
|
|
ctx_params.n_seq_max = 1;
|
|
ctx_params.no_perf = true;
|
|
|
|
llama_context *ctx = llama_init_from_model(model, ctx_params);
|
|
if (!ctx) die("failed to create context");
|
|
const int n_batch = std::min<int>((int)llama_n_batch(ctx), 2048);
|
|
llama_batch batch = llama_batch_init(n_batch, 0, 1);
|
|
|
|
std::ifstream mf(manifest_path, std::ios::binary);
|
|
if (!mf) {
|
|
std::fprintf(stderr, "open %s: %s\n", manifest_path, std::strerror(errno));
|
|
return 1;
|
|
}
|
|
std::ofstream out(out_path, std::ios::binary);
|
|
if (!out) {
|
|
std::fprintf(stderr, "open %s: %s\n", out_path, std::strerror(errno));
|
|
return 1;
|
|
}
|
|
out << "id\tprompt_tokens\ttarget_tokens\tnll\tavg_nll\tfirst_match\tgreedy_lcp\n";
|
|
|
|
std::string line;
|
|
int case_n = 0;
|
|
double total_nll = 0.0;
|
|
long total_tokens = 0;
|
|
long total_lcp = 0;
|
|
long first_matches = 0;
|
|
bool warned_template_fallback = false;
|
|
|
|
while (std::getline(mf, line)) {
|
|
strip_newline(line);
|
|
if (line.empty() || line[0] == '#') continue;
|
|
|
|
std::vector<std::string> cols = split_tab(line);
|
|
if (cols.size() < 3) die("bad manifest row");
|
|
const std::string &id = cols[0];
|
|
const std::string &prompt_path = cols[1];
|
|
const std::string &cont_path = cols[2];
|
|
|
|
std::string prompt_text = read_file(prompt_path.c_str());
|
|
std::string cont_text = read_file(cont_path.c_str());
|
|
|
|
std::string rendered;
|
|
bool used_template = false;
|
|
if (template_mode == "auto" && tmpl) {
|
|
rendered = render_template_prompt(tmpl, prompt_text, &used_template);
|
|
}
|
|
if (template_mode == "deepseek-ds4") {
|
|
rendered = render_deepseek_ds4_prompt(prompt_text);
|
|
used_template = true;
|
|
} else if (template_mode == "glm-ds4") {
|
|
rendered = render_glm_ds4_prompt(prompt_text);
|
|
used_template = true;
|
|
}
|
|
if (!used_template) {
|
|
if (template_mode == "auto" && !warned_template_fallback) {
|
|
std::fprintf(stderr,
|
|
"score_llama: llama.cpp chat template unavailable; "
|
|
"using DS4 GLM prompt fallback\n");
|
|
warned_template_fallback = true;
|
|
}
|
|
rendered = render_glm_ds4_prompt(prompt_text);
|
|
}
|
|
|
|
std::vector<llama_token> prompt =
|
|
tokenize(vocab, rendered, false, true);
|
|
std::vector<llama_token> target =
|
|
tokenize(vocab, cont_text, false, false);
|
|
|
|
if (prompt.empty()) die("empty prompt tokenization");
|
|
if (target.empty()) die("empty continuation tokenization");
|
|
if ((int)prompt.size() + (int)target.size() + 1 >= ctx_size) {
|
|
std::fprintf(stderr, "%s exceeds ctx=%d\n", id.c_str(), ctx_size);
|
|
return 1;
|
|
}
|
|
|
|
llama_memory_clear(llama_get_memory(ctx), true);
|
|
if (!decode_tokens(ctx, batch, prompt, 0, n_batch, true)) {
|
|
std::fprintf(stderr, "%s prompt decode failed\n", id.c_str());
|
|
return 1;
|
|
}
|
|
|
|
double nll = 0.0;
|
|
int lcp = 0;
|
|
bool still_matching = true;
|
|
bool first_match = false;
|
|
|
|
for (int i = 0; i < (int)target.size(); i++) {
|
|
const float *logits = llama_get_logits_ith(ctx, -1);
|
|
if (!logits) {
|
|
std::fprintf(stderr, "%s logits unavailable at target token %d\n",
|
|
id.c_str(), i);
|
|
return 1;
|
|
}
|
|
|
|
llama_token greedy = 0;
|
|
double lp = token_logprob(logits, n_vocab, target[(size_t)i], &greedy);
|
|
if (i == 0) first_match = (greedy == target[(size_t)i]);
|
|
if (still_matching && greedy == target[(size_t)i]) lcp++;
|
|
else still_matching = false;
|
|
nll += -lp;
|
|
|
|
if (!decode_chunk(ctx, batch, &target[(size_t)i], 1,
|
|
(int)prompt.size() + i, true)) {
|
|
std::fprintf(stderr, "%s target decode failed at token %d\n",
|
|
id.c_str(), i);
|
|
return 1;
|
|
}
|
|
}
|
|
|
|
const double avg = nll / (double)target.size();
|
|
out << id << '\t'
|
|
<< prompt.size() << '\t'
|
|
<< target.size() << '\t'
|
|
<< nll << '\t'
|
|
<< avg << '\t'
|
|
<< (first_match ? 1 : 0) << '\t'
|
|
<< lcp << '\n';
|
|
out.flush();
|
|
|
|
case_n++;
|
|
total_nll += nll;
|
|
total_tokens += (long)target.size();
|
|
total_lcp += lcp;
|
|
first_matches += first_match ? 1 : 0;
|
|
std::fprintf(stderr,
|
|
"%s cases=%d prompt=%zu target=%zu avg_nll=%.6f lcp=%d\n",
|
|
id.c_str(), case_n, prompt.size(), target.size(), avg, lcp);
|
|
}
|
|
|
|
std::fprintf(stderr,
|
|
"summary cases=%d tokens=%ld avg_nll=%.9f first_match=%ld avg_lcp=%.3f\n",
|
|
case_n,
|
|
total_tokens,
|
|
total_tokens ? total_nll / (double)total_tokens : 0.0,
|
|
first_matches,
|
|
case_n ? (double)total_lcp / (double)case_n : 0.0);
|
|
|
|
llama_batch_free(batch);
|
|
llama_free(ctx);
|
|
llama_model_free(model);
|
|
llama_backend_free();
|
|
return 0;
|
|
}
|