#include "common.h"

#include "llama.h"

#include <cmath>

#include <cstdio>

#include <cstring>

#include <ctime>

#include <sstream>

#include <thread>

#include <mutex>

#include <atomic>

#include <vector>

#include <array>

#include <fstream>

#include <sstream>

#if defined(_MSC_VER)

#pragma warning(disable: 4244 4267) // possible loss of data

#endif

struct results_perplexity {

std::vector<llama_token> tokens;

double ppl_value;

std::vector<float> logits;

std::vector<float> probs;

};

struct results_log_softmax {

double log_softmax;

float logit;

float prob;

};

static void write_logfile(

const llama_context * ctx, const gpt_params & params, const llama_model * model,

const struct results_perplexity & results

) {

if (params.logdir.empty()) {

return;

}

if (params.hellaswag) {

fprintf(stderr, "%s: warning: logging results is not implemented for HellaSwag. No files will be written.\n", __func__);

return;

}

const std::string timestamp = string_get_sortable_timestamp();

const bool success = fs_create_directory_with_parents(params.logdir);

if (!success) {

fprintf(stderr, "%s: warning: failed to create logdir %s, cannot write logfile\n",

__func__, params.logdir.c_str());

return;

}

const std::string logfile_path = params.logdir + timestamp + ".yml";

FILE * logfile = fopen(logfile_path.c_str(), "w");

if (logfile == NULL) {

fprintf(stderr, "%s: failed to open logfile %s\n", __func__, logfile_path.c_str());

return;

}

fprintf(logfile, "binary: main\n");

char model_desc[128];

llama_model_desc(model, model_desc, sizeof(model_desc));

yaml_dump_non_result_info(logfile, params, ctx, timestamp, results.tokens, model_desc);

fprintf(logfile, "\n");

fprintf(logfile, "######################\n");

fprintf(logfile, "# Perplexity Results #\n");

fprintf(logfile, "######################\n");

fprintf(logfile, "\n");

yaml_dump_vector_float(logfile, "logits", results.logits);

fprintf(logfile, "ppl_value: %f\n", results.ppl_value);

yaml_dump_vector_float(logfile, "probs", results.probs);

llama_dump_timing_info_yaml(logfile, ctx);

fclose(logfile);

}

static std::vector<float> softmax(const std::vector<float>& logits) {

std::vector<float> probs(logits.size());

float max_logit = logits[0];

for (float v : logits) {

max_logit = std::max(max_logit, v);

}

double sum_exp = 0.0;

for (size_t i = 0; i < logits.size(); i++) {

// Subtract the maximum logit value from the current logit value for numerical stability

const float logit = logits[i] - max_logit;

const float exp_logit = expf(logit);

sum_exp += exp_logit;

probs[i] = exp_logit;

}

for (size_t i = 0; i < probs.size(); i++) {

probs[i] /= sum_exp;

}

return probs;

}

static results_log_softmax log_softmax(int n_vocab, const float * logits, int tok) {

float max_logit = logits[0];

for (int i = 1; i < n_vocab; ++i) {

max_logit = std::max(max_logit, logits[i]);

}

double sum_exp = 0.0;

for (int i = 0; i < n_vocab; ++i) {

sum_exp += expf(logits[i] - max_logit);

}

return {logits[tok] - max_logit - log(sum_exp), logits[tok], expf(logits[tok] - max_logit) / (float) sum_exp};

}

static inline int nearest_int(float fval) {

//assert(fval <= 4194303.f);

float val = fval + 12582912.f;

int i; memcpy(&i, &val, sizeof(int));

return (i & 0x007fffff) - 0x00400000;

}

static double log_softmax(int n_vocab, const float * logits, uint16_t * log_prob, int tok) {

float max_logit = logits[0];

float min_logit = logits[0];

for (int i = 1; i < n_vocab; ++i) {

max_logit = std::max(max_logit, logits[i]);

min_logit = std::min(min_logit, logits[i]);

}

min_logit = std::max(min_logit, max_logit - 16);

double sum_exp = 0.0;

for (int i = 0; i < n_vocab; ++i) {

sum_exp += expf(logits[i] - max_logit);

}

const float log_sum_exp = log(sum_exp);

const float min_log_prob = min_logit - max_logit - log_sum_exp;

const float scale = (max_logit - min_logit)/65535.f;

float * d = (float *)log_prob;

d[0] = scale;

d[1] = min_log_prob;

log_prob += 4;

if (scale) {

const float inv_scale = 1/scale;

for (int i = 0; i < n_vocab; ++i) {

log_prob[i] = logits[i] > min_logit ? nearest_int(inv_scale*(logits[i] - min_logit)) : 0;

}

} else {

std::memset(log_prob, 0, n_vocab*sizeof(uint16_t));

}

return max_logit + log_sum_exp - logits[tok];

}

static void process_logits(

int n_vocab, const float * logits, const int * tokens, int n_token, std::vector<std::thread> & workers,

double & nll, double & nll2, float * logit_history, float * prob_history

) {

std::mutex mutex;

int counter = 0;

auto compute = [&mutex, &counter, &nll, &nll2, logit_history, prob_history, n_vocab, logits, tokens, n_token] () {

double local_nll = 0;

double local_nll2 = 0;

while (true) {

std::unique_lock<std::mutex> lock(mutex);

int i = counter++;

if (i >= n_token) {

nll += local_nll; nll2 += local_nll2;

break;

}

lock.unlock();

const results_log_softmax results = log_softmax(n_vocab, logits + i*n_vocab, tokens[i+1]);

const double v = -results.log_softmax;

local_nll += v;

local_nll2 += v*v;

logit_history[i] = results.logit;

prob_history[i] = results.prob;

}

};

for (auto & w : workers) {

w = std::thread(compute);

}

compute();

for (auto & w : workers) {

w.join();

}

}

static void process_logits(std::ostream& out, int n_vocab, const float * logits, const int * tokens, int n_token,

std::vector<std::thread> & workers, std::vector<uint16_t> & log_probs, double & nll, double & nll2) {

std::mutex mutex;

const int nv = 2*((n_vocab + 1)/2) + 4;

int counter = 0;

auto compute = [&mutex, &counter, &log_probs, &nll, &nll2, n_vocab, logits, tokens, n_token, nv] () {

double local_nll = 0;

double local_nll2 = 0;

while (true) {

std::unique_lock<std::mutex> lock(mutex);

int i = counter++;

if (i >= n_token) {

nll += local_nll; nll2 += local_nll2;

break;

}

lock.unlock();

const double v = log_softmax(n_vocab, logits + i*n_vocab, log_probs.data() + i*nv, tokens[i+1]);

local_nll += v;

local_nll2 += v*v;

}

};

for (auto & w : workers) {

w = std::thread(compute);

}

compute();

for (auto & w : workers) {

w.join();

}

out.write((const char *)log_probs.data(), n_token*nv*sizeof(uint16_t));

}

struct kl_divergence_result {

double sum_nll = 0.0;

double sum_nll2 = 0.0;

double sum_nll_base = 0.0;

double sum_nll_base2 = 0.0;

double sum_nll_nll_base = 0.0;

double sum_kld = 0.0;

double sum_kld2 = 0.0;

double sum_p_diff = 0.0;

double sum_p_diff2 = 0.0;

double sum_p_diff4 = 0.0;

float max_p_diff = 0.0f;

size_t n_same_top = 0.0;

size_t count = 0.0;

};

static std::pair<double, float> log_softmax(int n_vocab, const float * logits, const uint16_t * base_log_prob, int tok, kl_divergence_result & kld) {

float max_logit = logits[0];

int imax = 0;

for (int i = 1; i < n_vocab; ++i) {

if (logits[i] > max_logit) {

max_logit = logits[i];

imax = i;

}

}

double sum_exp = 0.0;

for (int i = 0; i < n_vocab; ++i) {

sum_exp += expf(logits[i] - max_logit);

}

const float log_sum_exp = log(sum_exp);

const float * d = (const float *)base_log_prob;

const float scale = d[0];

const float min_log_prob = d[1];

base_log_prob += 4;

const float nll = max_logit + log_sum_exp - logits[tok];

kld.sum_nll += nll;

kld.sum_nll2 += nll*nll;

const float nll_base = -(scale*base_log_prob[tok] + min_log_prob);

kld.sum_nll_base += nll_base;

kld.sum_nll_base2 += nll_base*nll_base;

kld.sum_nll_nll_base += nll*nll_base;

max_logit += log_sum_exp;

double sum = 0;

int imax_base = -1;

float p_log_base_max = 0;

for (int i = 0; i < n_vocab; ++i) {

const float p_log_base = scale*base_log_prob[i] + min_log_prob;

if (i == 0 || p_log_base > p_log_base_max) {

p_log_base_max = p_log_base;

imax_base = i;

}

if (p_log_base > -16.f) {

const float p_base = expf(p_log_base);

sum += p_base * (p_log_base - logits[i] + max_logit);

}

}

kld.sum_kld += sum;

kld.sum_kld2 += sum*sum;

++kld.count;

if (imax == imax_base) ++kld.n_same_top;

const float p_base = expf(-nll_base);

const float p = expf(-nll);

const float p_diff = p - p_base;

kld.sum_p_diff += p_diff;

const double p_diff2 = p_diff*p_diff;

kld.sum_p_diff2 += p_diff2;

kld.sum_p_diff4 += p_diff2*p_diff2;

kld.max_p_diff = std::max(kld.max_p_diff, std::fabs(p_diff));

return std::make_pair(sum, p_diff);

}

static void process_logits(int n_vocab, const float * logits, const int * tokens, int n_token,

std::vector<std::thread> & workers, const std::vector<uint16_t> & base_log_probs, kl_divergence_result & kld,

float * kld_values, float * p_diff_values) {

std::mutex mutex;

const int nv = 2*((n_vocab + 1)/2) + 4;

int counter = 0;

auto compute = [&mutex, &counter, &base_log_probs, &kld, n_vocab, logits, tokens, n_token, nv, kld_values, p_diff_values] () {

kl_divergence_result local_kld;

while (true) {

std::unique_lock<std::mutex> lock(mutex);

int i = counter++;

if (i >= n_token) {

kld.sum_nll += local_kld.sum_nll;

kld.sum_nll2 += local_kld.sum_nll2;

kld.sum_nll_base += local_kld.sum_nll_base;

kld.sum_nll_base2 += local_kld.sum_nll_base2;

kld.sum_nll_nll_base += local_kld.sum_nll_nll_base;

kld.sum_kld += local_kld.sum_kld;

kld.sum_kld2 += local_kld.sum_kld2;

kld.sum_p_diff += local_kld.sum_p_diff;

kld.sum_p_diff2 += local_kld.sum_p_diff2;

kld.sum_p_diff4 += local_kld.sum_p_diff4;

kld.n_same_top += local_kld.n_same_top;

kld.max_p_diff = std::max(kld.max_p_diff, local_kld.max_p_diff);

kld.count += local_kld.count;

break;

}

lock.unlock();

std::pair<double, float> v = log_softmax(n_vocab, logits + i*n_vocab, base_log_probs.data() + i*nv, tokens[i+1], local_kld);

kld_values[i] = (float)v.first;

p_diff_values[i] = v.second;

}

};

for (auto & w : workers) {

w = std::thread(compute);

}

compute();

for (auto & w : workers) {

w.join();

}

}

static results_perplexity perplexity_v2(llama_context * ctx, const gpt_params & params) {

// Download: https://huggingface.co/datasets/ggml-org/ci/resolve/main/wikitext-2-raw-v1.zip

// Run `./perplexity -m models/7B/ggml-model-q4_0.bin -f wiki.test.raw`

// Output: `perplexity: 13.5106 [114/114]`

// BOS tokens will be added for each chunk before eval

const bool add_bos = llama_should_add_bos_token(llama_get_model(ctx));

GGML_ASSERT(llama_add_eos_token(llama_get_model(ctx)) != 1);

fprintf(stderr, "%s: tokenizing the input ..\n", __func__);

std::vector<llama_token> tokens = ::llama_tokenize(ctx, params.prompt, true);

const int n_ctx = llama_n_ctx(ctx);

if (int(tokens.size()) < 2*n_ctx) {

fprintf(stderr, "%s: you need at least %d tokens to evaluate perplexity with a context of %d\n",__func__,2*n_ctx,

n_ctx);

fprintf(stderr, "%s: the data file you provided tokenizes to only %zu tokens\n",__func__,tokens.size());

return {std::move(tokens), 0., {}, {}};

}

std::vector<float> logit_history;

std::vector<float> prob_history;

logit_history.resize(tokens.size());

prob_history.resize(tokens.size());

if (params.ppl_stride <= 0) {

fprintf(stderr, "%s: stride is %d but must be greater than zero!\n",__func__,params.ppl_stride);

return {tokens, -1, logit_history, prob_history};

}

const int calc_chunk = n_ctx;

fprintf(stderr, "%s: have %zu tokens. Calculation chunk = %d\n", __func__, tokens.size(), calc_chunk);

if (int(tokens.size()) <= calc_chunk) {

fprintf(stderr, "%s: there are only %zu tokens, this is not enough for a context size of %d and stride %d\n",__func__,

tokens.size(), n_ctx, params.ppl_stride);

return {tokens, -1, logit_history, prob_history};

}

const int n_chunk_max = (tokens.size() - calc_chunk + params.ppl_stride - 1) / params.ppl_stride;

const int n_chunk = params.n_chunks < 0 ? n_chunk_max : std::min(params.n_chunks, n_chunk_max);

const int n_vocab = llama_n_vocab(llama_get_model(ctx));

const int n_batch = params.n_batch;

int count = 0;

double nll = 0.0;

fprintf(stderr, "%s: calculating perplexity over %d chunks, batch_size=%d\n", __func__, n_chunk, n_batch);

for (int i = 0; i < n_chunk; ++i) {

const int start = i * params.ppl_stride;

const int end = start + calc_chunk;

const int num_batches = (calc_chunk + n_batch - 1) / n_batch;

//fprintf(stderr, "%s: evaluating %d...%d using %d batches\n", __func__, start, end, num_batches);

std::vector<float> logits;

const auto t_start = std::chrono::high_resolution_clock::now();

// clear the KV cache

llama_kv_cache_clear(ctx);

for (int j = 0; j < num_batches; ++j) {

const int batch_start = start + j * n_batch;

const int batch_size = std::min(end - batch_start, n_batch);

//fprintf(stderr, " Batch %d: starts at %d, size is %d, n_past is %d\n",j,batch_start,batch_size,j * n_batch);

// TODO: use llama_batch.logits instead of relying on logits_all == true

if (llama_decode(ctx, llama_batch_get_one(tokens.data() + batch_start, batch_size, j * n_batch, 0))) {

//fprintf(stderr, "%s : failed to eval\n", __func__);

return {tokens, -1, logit_history, prob_history};

}

// save original token and restore it after eval

const auto token_org = tokens[batch_start];

// add BOS token for the first batch of each chunk

if (add_bos && j == 0) {

tokens[batch_start] = llama_token_bos(llama_get_model(ctx));

}

const auto batch_logits = llama_get_logits(ctx);

logits.insert(logits.end(), batch_logits, batch_logits + batch_size * n_vocab);

if (j == 0) {

tokens[batch_start] = token_org;

}

}

const auto t_end = std::chrono::high_resolution_clock::now();

if (i == 0) {

const float t_total = std::chrono::duration<float>(t_end - t_start).count();

fprintf(stderr, "%s: %.2f seconds per pass - ETA ", __func__, t_total);

int total_seconds = (int)(t_total * n_chunk);

if (total_seconds >= 60*60) {

fprintf(stderr, "%d hours ", total_seconds / (60*60));

total_seconds = total_seconds % (60*60);

}

fprintf(stderr, "%.2f minutes\n", total_seconds / 60.0);

}

//fprintf(stderr, "%s: using tokens %d...%d\n",__func__,params.n_ctx - params.ppl_stride + start, params.n_ctx + start);

for (int j = n_ctx - params.ppl_stride - 1; j < n_ctx - 1; ++j) {

// Calculate probability of next token, given the previous ones.

const std::vector<float> tok_logits(

logits.begin() + (j + 0) * n_vocab,

logits.begin() + (j + 1) * n_vocab);

const float prob = softmax(tok_logits)[tokens[start + j + 1]];

logit_history[start + j + 1] = tok_logits[tokens[start + j + 1]];

prob_history[start + j + 1] = prob;

nll += -std::log(prob);

++count;

}

// perplexity is e^(average negative log-likelihood)

if (params.ppl_output_type == 0) {

printf("[%d]%.4lf,", i + 1, std::exp(nll / count));

} else {

printf("%8d %.4lf\n", i*params.ppl_stride, std::exp(nll / count));

}

fflush(stdout);

}

printf("\n");

return {tokens, std::exp(nll / count), logit_history, prob_history};

}

static results_perplexity perplexity(llama_context * ctx, const gpt_params & params, const int32_t n_ctx) {

if (params.ppl_stride > 0) {

return perplexity_v2(ctx, params);

}

// Download: https://huggingface.co/datasets/ggml-org/ci/resolve/main/wikitext-2-raw-v1.zip

// Run `./llama-perplexity -m models/7B/ggml-model-q4_0.bin -f wiki.test.raw`

// Output: `perplexity: 13.5106 [114/114]`

// BOS tokens will be added for each chunk before eval

const bool add_bos = llama_should_add_bos_token(llama_get_model(ctx));

GGML_ASSERT(llama_add_eos_token(llama_get_model(ctx)) != 1);

std::ofstream logits_stream;

if (!params.logits_file.empty()) {

logits_stream.open(params.logits_file.c_str(), std::ios::binary);

if (!logits_stream.is_open()) {

fprintf(stderr, "%s: failed to open %s for writing\n", __func__, params.logits_file.c_str());

return {};

}

fprintf(stderr, "%s: saving all logits to %s\n", __func__, params.logits_file.c_str());

logits_stream.write("_logits_", 8);

logits_stream.write(reinterpret_cast<const char *>(&n_ctx), sizeof(n_ctx));

}

auto tim1 = std::chrono::high_resolution_clock::now();

fprintf(stderr, "%s: tokenizing the input ..\n", __func__);

std::vector<llama_token> tokens = ::llama_tokenize(ctx, params.prompt, true);

auto tim2 = std::chrono::high_resolution_clock::now();

fprintf(stderr, "%s: tokenization took %g ms\n",__func__,1e-3*std::chrono::duration_cast<std::chrono::microseconds>(tim2-tim1).count());

if (int(tokens.size()) < 2*n_ctx) {

fprintf(stderr, "%s: you need at least %d tokens to evaluate perplexity with a context of %d\n",__func__,2*n_ctx,

n_ctx);

fprintf(stderr, "%s: the data file you provided tokenizes to only %zu tokens\n",__func__,tokens.size());

return {std::move(tokens), 0., {}, {}};

}

std::vector<float> logit_history;

logit_history.resize(tokens.size());

std::vector<float> prob_history;

prob_history.resize(tokens.size());

const int n_chunk_max = tokens.size() / n_ctx;

const int n_chunk = params.n_chunks < 0 ? n_chunk_max : std::min(params.n_chunks, n_chunk_max);

const int n_vocab = llama_n_vocab(llama_get_model(ctx));

const int n_batch = params.n_batch;

int count = 0;

double nll = 0.0;

double nll2 = 0.0;

const int num_batches = (n_ctx + n_batch - 1) / n_batch;

const int n_seq = std::max(1, n_batch / n_ctx);

GGML_ASSERT(n_batch < n_ctx || n_batch % n_ctx == 0);

GGML_ASSERT(params.n_ctx == n_seq * n_ctx);

llama_batch batch = llama_batch_init(std::min(n_batch, n_ctx*n_seq), 0, 1);

std::vector<float> logits;

if (num_batches > 1) {

logits.reserve((size_t)n_ctx * n_vocab);

}

fprintf(stderr, "%s: calculating perplexity over %d chunks, n_ctx=%d, batch_size=%d, n_seq=%d\n", __func__, n_chunk, n_ctx, n_batch, n_seq);

std::vector<std::thread> workers(std::thread::hardware_concurrency() - 1);

std::vector<uint16_t> log_probs;

if (!params.logits_file.empty()) {

logits_stream.write((const char *)&n_vocab, sizeof(n_vocab));

logits_stream.write((const char *)&n_chunk, sizeof(n_chunk));

logits_stream.write((const char *)tokens.data(), n_chunk*n_ctx*sizeof(tokens[0]));

const int nv = 2*((n_vocab + 1)/2) + 4;

log_probs.resize(n_ctx * nv);

}

// We get the logits for all the tokens in the context window (params.n_ctx)

// from llama_eval above. Now, based on https://huggingface.co/docs/transformers/perplexity,

// calculate the perplexity over the last half of the window (so the model always has

// some context to predict the token).

//

// We rely on the fact that attention in the forward pass only looks at previous

// tokens here, so the logits returned for each token are an accurate representation

// of what the model would have predicted at that point.

//

// Example, we have a context window of 512, we will compute perplexity for each of the

// last 256 tokens. Then, we split the input up into context window size chunks to

// process the entire prompt.

const int first = n_ctx/2;

for (int i = 0; i < n_chunk; i += n_seq) {

const int start = i * n_ctx;

const int end = start + n_ctx;

const int n_seq_batch = std::min(n_seq, n_chunk - i);

const auto t_start = std::chrono::high_resolution_clock::now();

// clear the KV cache

llama_kv_cache_clear(ctx);

for (int j = 0; j < num_batches; ++j) {

const int batch_start = start + j * n_batch;

const int batch_size = std::min(end - batch_start, n_batch);

int n_outputs = 0;

batch.n_tokens = 0;

for (int seq = 0; seq < n_seq_batch; seq++) {

int seq_start = batch_start + seq*n_ctx;

// save original token and restore it after eval

const auto token_org = tokens[seq_start];

// add BOS token for the first batch of each chunk

if (add_bos && j == 0) {

tokens[seq_start] = llama_token_bos(llama_get_model(ctx));

}

for (int k = 0; k < batch_size; ++k) {

const int idx = seq*n_ctx + k;

batch.token [idx] = tokens[seq_start + k];

batch.pos [idx] = j*n_batch + k;

batch.n_seq_id[idx] = 1;

batch.seq_id [idx][0] = seq;

batch.logits [idx] = batch.pos[idx] >= first ? 1 : 0;

n_outputs += batch.logits[idx] != 0;

}

batch.n_tokens += batch_size;

// restore the original token in case it was set to BOS

tokens[seq_start] = token_org;

}

if (llama_decode(ctx, batch)) {

fprintf(stderr, "%s : failed to eval\n", __func__);

return {tokens, -1, logit_history, prob_history};

}

if (num_batches > 1 && n_outputs > 0) {

const auto * batch_logits = llama_get_logits(ctx);

logits.insert(logits.end(), batch_logits, batch_logits + n_outputs * n_vocab);

}

}

if (i == 0) {

llama_synchronize(ctx);

const auto t_end = std::chrono::high_resolution_clock::now();

const float t_total = std::chrono::duration<float>(t_end - t_start).count();

fprintf(stderr, "%s: %.2f seconds per pass - ETA ", __func__, t_total);

int total_seconds = (int)(t_total*n_chunk/n_seq);

if (total_seconds >= 60*60) {

fprintf(stderr, "%d hours ", total_seconds / (60*60));

total_seconds = total_seconds % (60*60);

}

fprintf(stderr, "%.2f minutes\n", total_seconds / 60.0);

}

for (int seq = 0; seq < n_seq_batch; seq++) {

const float * all_logits = num_batches > 1 ? logits.data() : llama_get_logits_ith(ctx, seq*n_ctx + first);

llama_token * tokens_data = tokens.data() + start + seq*n_ctx + first;

if (!params.logits_file.empty()) {

process_logits(logits_stream, n_vocab, all_logits,

tokens_data, n_ctx - 1 - first,

workers, log_probs, nll, nll2);

} else {

process_logits(n_vocab, all_logits,

tokens_data, n_ctx - 1 - first,

workers, nll, nll2,

logit_history.data() + start + seq*n_ctx + first,

prob_history.data() + start + seq*n_ctx + first);

}

count += n_ctx - first - 1;

// perplexity is e^(average negative log-likelihood)

if (params.ppl_output_type == 0) {

printf("[%d]%.4lf,", i + seq + 1, std::exp(nll / count));

} else {

double av = nll/count;

double av2 = nll2/count - av*av;

if (av2 > 0) av2 = sqrt(av2/(count-1));

printf("%8d %.4lf %4lf %4lf\n", i*n_ctx, std::exp(nll / count), av, av2);

}

}

fflush(stdout);

logits.clear();

}

printf("\n");

nll2 /= count;

nll /= count;

const double ppl = exp(nll);

nll2 -= nll * nll;

if (nll2 > 0) {

nll2 = sqrt(nll2/(count-1));

printf("Final estimate: PPL = %.4lf +/- %.5lf\n", ppl, nll2*ppl);

} else {

printf("Unexpected negative standard deviation of log(prob)\n");

}

llama_batch_free(batch);

return {tokens, ppl, logit_history, prob_history};

}

static bool decode_helper(llama_context * ctx, llama_batch & batch, std::vector<float> & batch_logits, int32_t n_batch, int32_t n_vocab) {

int prev_outputs = 0;

for (int32_t i = 0; i < (int32_t) batch.n_tokens; i += n_batch) {

const int32_t n_tokens = std::min(n_batch, (int32_t) (batch.n_tokens - i));

llama_batch batch_view = {

n_tokens,

batch.token + i,

nullptr,

batch.pos + i,

batch.n_seq_id + i,

batch.seq_id + i,

batch.logits + i,

0, 0, 0, // unused

};

const int ret = llama_decode(ctx, batch_view);

if (ret != 0) {

LOG_TEE("failed to decode the batch, n_batch = %d, ret = %d\n", n_batch, ret);

return false;

}

int n_outputs = 0;

for (int i = 0; i < n_tokens; ++i) {

n_outputs += batch_view.logits[i] != 0;

}

memcpy(batch_logits.data() + prev_outputs*n_vocab, llama_get_logits(ctx), n_outputs*n_vocab*sizeof(float));

prev_outputs += n_outputs;

}

return true;

}

#define K_TOKEN_CHUNK 4

static void compute_logprobs(const float * batch_logits, int n_vocab, std::vector<std::thread>& workers,

const std::vector<std::pair<size_t, llama_token>>& eval_pairs, std::vector<float>& eval_results) {

if (eval_results.size() != eval_pairs.size()) {

eval_results.resize(eval_pairs.size());

}

if (eval_pairs.empty()) return;

size_t max_threads = std::min((eval_pairs.size() + K_TOKEN_CHUNK - 1)/K_TOKEN_CHUNK, workers.size());

std::atomic<int> counter(0);

auto compute = [&counter, &eval_pairs, &eval_results, batch_logits, n_vocab] () {

float local_logprobs[K_TOKEN_CHUNK];

while (true) {

size_t first = counter.fetch_add(K_TOKEN_CHUNK, std::memory_order_relaxed);

if (first >= eval_results.size()) break;

size_t last = std::min(first + K_TOKEN_CHUNK, eval_results.size());

for (size_t i = first; i < last; ++i) {

auto logits = batch_logits + eval_pairs[i].first * n_vocab;

float max_logit = logits[0];

for (int j = 1; j < n_vocab; ++j) {

max_logit = std::max(max_logit, logits[j]);

}

float sum_p = 0.f;

for (int j = 0; j < n_vocab; ++j) {

sum_p += expf(logits[j] - max_logit);

}

local_logprobs[i - first] = logits[eval_pairs[i].second] - max_logit - std::log(sum_p);

}

std::memcpy(eval_results.data() + first, local_logprobs, (last - first)*sizeof(float));

}

};

for (size_t it = 0; it < max_threads; ++it) {

workers[it] = std::thread(compute);

}

for (size_t it = 0; it < max_threads; ++it) {

workers[it].join();

}

}

static void hellaswag_score(llama_context * ctx, const gpt_params & params) {

// Calculates hellaswag score (acc_norm) from prompt

//

// Data extracted from the HellaSwag validation dataset (MIT license) https://github.com/rowanz/hellaswag/blob/master/data/hellaswag_val.jsonl

// All used data fields are preprocessed as in https://github.com/EleutherAI/lm-evaluation-harness/blob/df3da98c5405deafd519c2ddca52bb7c3fe36bef/lm_eval/tasks/hellaswag.py#L62-L68

//

// All 10042 tasks should be extracted to keep the results standardized like other implementations.

//

// Datafile layout:

// ['??'] denotes json fields

// 6 lines per task:

// ['activity_label'] + ": " +['ctx'] - The first part of the query, the context

// ['label'] - The index the best common sense ending aka gold ending

// ['endings'][0] - Endings added to the first part of the query

// ['endings'][1]

// ['endings'][2]

// ['endings'][3]

std::vector<std::string> prompt_lines;

std::istringstream strstream(params.prompt);

std::string line;

while (std::getline(strstream,line,'\n')) {

prompt_lines.push_back(line);

}

if (prompt_lines.size() % 6 != 0) {

fprintf(stderr, "%s : number of lines in prompt not a multiple of 6.\n", __func__);

return;

}

size_t hs_task_count = prompt_lines.size()/6;

fprintf(stderr, "%s : loaded %zu tasks from prompt.\n", __func__, hs_task_count);

const bool is_spm = llama_vocab_type(llama_get_model(ctx)) == LLAMA_VOCAB_TYPE_SPM;

fprintf(stderr, "================================= is_spm = %d\n", is_spm);

// The tasks should be randomized so the score stabilizes quickly.

bool randomize_tasks = true;

// Number of tasks to use when computing the score

if (params.hellaswag_tasks < hs_task_count) {

hs_task_count = params.hellaswag_tasks;

}

// The random seed should not impact the final result if the computation is done over enough tasks, so kept hardcoded for now

std::mt19937 rng(1);

// Dataholder for hellaswag tasks

struct hs_data_t {

std::string context;

size_t gold_ending_idx;

std::string ending[4];

size_t ending_logprob_count[4];

double ending_logprob[4];

size_t i_logits; // starting index of logits in the llama_batch

size_t common_prefix; // max number of initial tokens that are the same in all sentences

size_t required_tokens; // needed number of tokens to evaluate all 4 endings

std::vector<llama_token> seq_tokens[4];

};

fprintf(stderr, "%s : selecting %zu %s tasks.\n", __func__, hs_task_count, (randomize_tasks?"randomized":"the first") );

// Select and read data from prompt lines

std::vector<hs_data_t> hs_data(hs_task_count);

for (size_t i = 0; i < hs_task_count; i++) {

size_t idx = i;

auto & hs_cur = hs_data[i];

// Select a random example of those left in the prompt

if (randomize_tasks) {

std::uniform_int_distribution<size_t> dist(0, prompt_lines.size()/6-1 ) ;

idx = dist(rng);

}

hs_cur.context = prompt_lines[idx*6];

hs_cur.gold_ending_idx = std::stoi( prompt_lines[idx*6+1] );

for (size_t j = 0; j < 4; j++) {

hs_cur.ending[j] = prompt_lines[idx*6+2+j];

hs_cur.seq_tokens[j] = ::llama_tokenize(ctx, hs_cur.context + " " + hs_cur.ending[j], true);

}

// determine the common prefix of the endings

hs_cur.common_prefix = 0;

for (size_t k = 0; k < hs_cur.seq_tokens[0].size(); k++) {

if (hs_cur.seq_tokens[0][k] != hs_cur.seq_tokens[1][k] ||

hs_cur.seq_tokens[0][k] != hs_cur.seq_tokens[2][k] ||

hs_cur.seq_tokens[0][k] != hs_cur.seq_tokens[3][k]) {

break;

}

hs_cur.common_prefix++;

}

hs_cur.required_tokens = hs_cur.common_prefix +

hs_cur.seq_tokens[0].size() - hs_cur.common_prefix +

hs_cur.seq_tokens[1].size() - hs_cur.common_prefix +

hs_cur.seq_tokens[2].size() - hs_cur.common_prefix +

hs_cur.seq_tokens[3].size() - hs_cur.common_prefix;

//GGML_ASSERT(hs_cur.common_prefix >= ::llama_tokenize(ctx, hs_cur.context, true).size());

// Delete the selected random example from the prompt

if (randomize_tasks) {

prompt_lines.erase( std::next(prompt_lines.begin(),idx*6) , std::next(prompt_lines.begin(),idx*6+6) );

}

}

fprintf(stderr, "%s : calculating hellaswag score over selected tasks.\n", __func__);

printf("\ntask\tacc_norm\n");

double acc = 0.0f;

const int n_vocab = llama_n_vocab(llama_get_model(ctx));

const int n_ctx = llama_n_ctx(ctx);

const int n_batch = params.n_batch;

const int max_tasks_per_batch = 32;

const int max_seq = std::min(4*max_tasks_per_batch, (int) llama_n_seq_max(ctx));

llama_batch batch = llama_batch_init(n_ctx, 0, 4);

std::vector<float> tok_logits(n_vocab);

// TODO: this could be made smaller; it's currently the worst-case size

std::vector<float> batch_logits(n_vocab*n_ctx);

std::vector<std::pair<size_t, llama_token>> eval_pairs;

std::vector<float> eval_results;

std::vector<std::thread> workers(std::thread::hardware_concurrency());

for (size_t i0 = 0; i0 < hs_task_count; i0++) {

int n_cur = 0;

size_t i1 = i0;

size_t i_logits = 0; // this tells us how many logits were needed before this point in the batch

llama_batch_clear(batch);

// batch as much tasks as possible into the available context

// each task has 4 unique sequence ids - one for each ending

// the common prefix is shared among the 4 sequences to save tokens

// we extract logits only from the last common token and from all ending tokens of each sequence

while (n_cur + (int) hs_data[i1].required_tokens <= n_ctx) {

auto & hs_cur = hs_data[i1];

int n_logits = 0;

const int s0 = 4*(i1 - i0);

if (s0 + 4 > max_seq) {

break;

}

for (size_t i = 0; i < hs_cur.common_prefix; ++i) {

llama_batch_add(batch, hs_cur.seq_tokens[0][i], i, { s0 + 0, s0 + 1, s0 + 2, s0 + 3 }, false);

}

batch.logits[batch.n_tokens - 1] = true; // we need logits for the last token of the common prefix

n_logits += 1;

for (int s = 0; s < 4; ++s) {

const size_t seq_tokens_size = hs_cur.seq_tokens[s].size();

// TODO: don't evaluate the last token of each sequence

for (size_t i = hs_cur.common_prefix; i < seq_tokens_size; ++i) {

const bool needs_logits = i < seq_tokens_size - 1;

llama_batch_add(batch, hs_cur.seq_tokens[s][i], i, { s0 + s }, needs_logits);

n_logits += needs_logits;

}

}

hs_cur.i_logits = i_logits;

i_logits += n_logits;

n_cur += hs_data[i1].required_tokens;

if (++i1 == hs_task_count) {

break;

}

}

if (i0 == i1) {

fprintf(stderr, "%s : task %zu does not fit in the context window\n", __func__, i0);

return;

}

llama_kv_cache_clear(ctx);

// decode all tasks [i0, i1)

if (!decode_helper(ctx, batch, batch_logits, n_batch, n_vocab)) {

fprintf(stderr, "%s: llama_decode() failed\n", __func__);

return;

}

// Compute log-probs in parallel

// First we collect all tasks

eval_pairs.clear();

for (size_t i = i0; i < i1; ++i) {

auto & hs_cur = hs_data[i];

size_t li = 1; // skip the last logit of the common prefix (computed separately below)

for (int s = 0; s < 4; ++s) {

for (size_t j = hs_cur.common_prefix; j < hs_cur.seq_tokens[s].size() - 1; j++) {

eval_pairs.emplace_back(hs_cur.i_logits + li++, hs_cur.seq_tokens[s][j + 1]);

}

}

}

// Then we do the actual calculation

compute_logprobs(batch_logits.data(), n_vocab, workers, eval_pairs, eval_results);

size_t ir = 0;

// compute the logprobs for each ending of the decoded tasks

for (size_t i = i0; i < i1; ++i) {

auto & hs_cur = hs_data[i];

// get the logits of the last token of the common prefix

std::memcpy(tok_logits.data(), batch_logits.data() + n_vocab*hs_cur.i_logits, n_vocab*sizeof(float));

const auto first_probs = softmax(tok_logits);

for (int s = 0; s < 4; ++s) {

hs_cur.ending_logprob_count[s] = 1;

hs_cur.ending_logprob[s] = std::log(first_probs[hs_cur.seq_tokens[s][hs_cur.common_prefix]]);

for (size_t j = hs_cur.common_prefix; j < hs_cur.seq_tokens[s].size() - 1; j++) {

hs_cur.ending_logprob[s] += eval_results[ir++];

hs_cur.ending_logprob_count[s]++;

}

hs_cur.ending_logprob[s] /= hs_cur.ending_logprob_count[s];

}

// Find the ending with maximum logprob

size_t ending_logprob_max_idx = 0;

double ending_logprob_max_val = hs_cur.ending_logprob[0];

for (size_t s = 1; s < 4; s++) {

if (hs_cur.ending_logprob[s] > ending_logprob_max_val) {

ending_logprob_max_idx = s;

ending_logprob_max_val = hs_cur.ending_logprob[s];

}

}