From 02f0c6fe7f9b7be24c7d339aed016e54a92388ea Mon Sep 17 00:00:00 2001 From: beiller Date: Sun, 12 Mar 2023 16:23:15 -0400 Subject: Add back top_k (#56) * Add back top_k * Update utils.cpp * Update utils.h --------- Co-authored-by: Bill Hamilton Co-authored-by: Georgi Gerganov --- main.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) (limited to 'main.cpp') diff --git a/main.cpp b/main.cpp index a11d755a..01556143 100644 --- a/main.cpp +++ b/main.cpp @@ -825,6 +825,7 @@ int main(int argc, char ** argv) { if (i >= embd_inp.size()) { // sample next token + const float top_k = params.top_k; const float top_p = params.top_p; const float temp = params.temp; const float repeat_penalty = params.repeat_penalty; @@ -836,7 +837,7 @@ int main(int argc, char ** argv) { { const int64_t t_start_sample_us = ggml_time_us(); - id = llama_sample_top_p(vocab, logits.data() + (logits.size() - n_vocab), last_n_tokens, repeat_penalty, top_p, temp, rng); + id = llama_sample_top_p_top_k(vocab, logits.data() + (logits.size() - n_vocab), last_n_tokens, repeat_penalty, top_k, top_p, temp, rng); last_n_tokens.erase(last_n_tokens.begin()); last_n_tokens.push_back(id); -- cgit v1.2.3