diff --git a/common/sampling.cpp b/common/sampling.cpp index 5698c0263..ba5504ed0 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -383,7 +383,7 @@ struct common_sampler * common_sampler_init( samplers.push_back(llama_sampler_init_infill(vocab)); break; case COMMON_SAMPLER_TYPE_PENALTIES: - samplers.push_back(llama_sampler_init_penalties(params.penalty_last_n, params.penalty_repeat, params.penalty_freq, params.penalty_present)); + samplers.push_back(llama_sampler_init_penalties(llama_vocab_n_tokens(vocab), params.penalty_last_n, params.penalty_repeat, params.penalty_freq, params.penalty_present)); break; case COMMON_SAMPLER_TYPE_ADAPTIVE_P: // the `adaptive-p` sampler is like `dist` and `mirostat` in that it selects diff --git a/include/llama.h b/include/llama.h index f2d7e3885..fb2ca38ce 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1256,7 +1256,6 @@ extern "C" { struct ggml_tensor * probs; struct ggml_tensor * sampled; struct ggml_tensor * candidates; - int64_t n_vocab; }; // user code can implement the interface below in order to create custom llama_sampler @@ -1425,6 +1424,7 @@ extern "C" { /// NOTE: Avoid using on the full vocabulary as searching for repeated tokens can become slow. For example, apply top-k or top-p sampling first. LLAMA_API struct llama_sampler * llama_sampler_init_penalties( + int32_t n_vocab, int32_t penalty_last_n, // last n tokens to penalize (0 = disable penalty, -1 = context size) float penalty_repeat, // must be > 0.0, 1.0 = disabled float penalty_freq, // must be finite, 0.0 = disabled diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 9dde345df..2be3b75fb 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3683,7 +3683,6 @@ void llm_graph_context::build_sampling() const { /*.probs =*/ nullptr, /*.sampled =*/ nullptr, /*.candidates =*/ nullptr, - /*.n_vocab =*/ logits_seq->ne[0], }; assert(sampler->iface->backend_apply); diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index b2f1abe73..6cf2d27cf 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -589,7 +589,6 @@ static bool llama_sampler_backend_support( /*.probs = */ nullptr, /*.sampled = */ nullptr, /*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n), - /*.n_vocab = */ n, }; ggml_cgraph * gf = ggml_new_graph(ctx); @@ -2640,6 +2639,7 @@ struct llama_sampler * llama_sampler_init_grammar_lazy_patterns( // penalties struct llama_sampler_penalties : public llama_sampler_backend { + const int32_t n_vocab; const int32_t penalty_last_n; const float penalty_repeat; const float penalty_freq; @@ -2655,7 +2655,6 @@ struct llama_sampler_penalties : public llama_sampler_backend { ggml_tensor * inp_counts = nullptr; // backend helpers - int32_t n_vocab = 0; int32_t n_max = 0; bool has_candidates = false; @@ -2676,11 +2675,13 @@ struct llama_sampler_penalties : public llama_sampler_backend { } llama_sampler_penalties( + int32_t n_vocab, int32_t penalty_last_n, float penalty_repeat, float penalty_freq, float penalty_present) : llama_sampler_backend("penalties") + , n_vocab (n_vocab) , penalty_last_n (penalty_last_n) , penalty_repeat (penalty_repeat) , penalty_freq (penalty_freq) @@ -2766,6 +2767,7 @@ static void llama_sampler_penalties_reset(struct llama_sampler * smpl) { static struct llama_sampler * llama_sampler_penalties_clone(const struct llama_sampler * smpl) { const auto * ctx = (const llama_sampler_penalties *) smpl->ctx; auto * result = llama_sampler_init_penalties( + ctx->n_vocab, ctx->penalty_last_n, ctx->penalty_repeat, ctx->penalty_freq, @@ -2811,10 +2813,9 @@ static void llama_sampler_penalties_backend_apply( return; } - GGML_ASSERT(data->n_vocab > 0 && data->n_vocab <= INT32_MAX); + GGML_ASSERT(sctx->n_vocab > 0); sctx->has_candidates = data->candidates != nullptr; - sctx->n_vocab = (int32_t) data->n_vocab; sctx->n_max = std::min(sctx->penalty_last_n, sctx->n_vocab); sctx->inp_token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sctx->n_max); @@ -2965,6 +2966,7 @@ static struct llama_sampler_i llama_sampler_penalties_i = { }; struct llama_sampler * llama_sampler_init_penalties( + int32_t n_vocab, int32_t penalty_last_n, float penalty_repeat, float penalty_freq, @@ -2979,6 +2981,7 @@ struct llama_sampler * llama_sampler_init_penalties( return llama_sampler_init( /* .iface = */ &llama_sampler_penalties_i, /* .ctx = */ new llama_sampler_penalties( + n_vocab, penalty_last_n, penalty_repeat, penalty_freq, diff --git a/tests/test-backend-sampler.cpp b/tests/test-backend-sampler.cpp index 1a46468ba..1165f46f0 100644 --- a/tests/test-backend-sampler.cpp +++ b/tests/test-backend-sampler.cpp @@ -823,6 +823,7 @@ enum class penalties_position { static void add_filter_and_penalties( llama_sampler * chain, const sampler_init_fn & init_filter, + int32_t n_vocab, int32_t penalty_last_n, float penalty_repeat, float penalty_freq, @@ -830,7 +831,7 @@ static void add_filter_and_penalties( penalties_position position) { const auto add_penalties = [&]() { llama_sampler_chain_add(chain, llama_sampler_init_penalties( - penalty_last_n, penalty_repeat, penalty_freq, penalty_present)); + n_vocab, penalty_last_n, penalty_repeat, penalty_freq, penalty_present)); }; if (position == penalties_position::before_filter) { @@ -1006,7 +1007,7 @@ static sampler_comparison_output run_penalties_comparison( const std::vector raw_logits = decode_raw_logits(params, prompt); const auto add_samplers = [&](llama_sampler * chain) { llama_sampler_chain_add(chain, llama_sampler_init_penalties( - penalty_last_n, penalty_repeat, penalty_freq, penalty_present)); + llama_vocab_n_tokens(vocab), penalty_last_n, penalty_repeat, penalty_freq, penalty_present)); }; const auto accept_history = [&](llama_sampler * chain) { accept_prompt(chain, vocab, prompt); @@ -1105,7 +1106,7 @@ static void compare_top_k_penalties_logits( GGML_ASSERT(excluded_history_token != LLAMA_TOKEN_NULL); const auto add_samplers = [&](llama_sampler * chain) { - add_filter_and_penalties(chain, init_top_k, + add_filter_and_penalties(chain, init_top_k, n_vocab, penalty_last_n, penalty_repeat, penalty_freq, penalty_present, position); }; @@ -1190,7 +1191,7 @@ static void compare_masking_penalties_logits( GGML_ASSERT(masked_token != LLAMA_TOKEN_NULL); const auto add_samplers = [&](llama_sampler * chain) { - add_filter_and_penalties(chain, init_filter, + add_filter_and_penalties(chain, init_filter, n_vocab, penalty_last_n, penalty_repeat, penalty_freq, penalty_present, position); }; auto accept_history = [&](llama_sampler * smpl) { @@ -1218,7 +1219,7 @@ static void compare_masking_penalties_logits( GGML_ASSERT(fabsf(expected_logits.at(penalized_token) - raw_logits[penalized_token]) > 1e-6f); } else { llama_sampler_ptr penalties(llama_sampler_init_penalties( - penalty_last_n, penalty_repeat, penalty_freq, penalty_present)); + n_vocab, penalty_last_n, penalty_repeat, penalty_freq, penalty_present)); accept_history(penalties.get()); const std::unordered_map penalized_logits = map_logits(apply_cpu_sampler(raw_logits, penalties.get())); diff --git a/tests/test-sampling.cpp b/tests/test-sampling.cpp index 2aecff90e..297f76015 100644 --- a/tests/test-sampling.cpp +++ b/tests/test-sampling.cpp @@ -144,7 +144,7 @@ static void test_penalties( sampler_tester tester(probs, probs_expected); - auto * sampler = llama_sampler_init_penalties(last_tokens.size(), repeat_penalty, alpha_frequency, alpha_presence); + auto * sampler = llama_sampler_init_penalties((int32_t) probs.size(), (int32_t) last_tokens.size(), repeat_penalty, alpha_frequency, alpha_presence); for (size_t i = 0; i < last_tokens.size(); i++) { llama_sampler_accept(sampler, last_tokens[i]);