Fix to guarantee K-Shift on the first step only

This commit is contained in:
MaggotHATE 2024-11-08 11:55:10 +05:00
parent 840a2b1c61
commit 877a495245

View file

@ -1107,14 +1107,18 @@ static const char * llama_sampler_k_shift_name(const struct llama_sampler * /*sm
static void llama_sampler_k_shift_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) { static void llama_sampler_k_shift_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
auto * ctx = (llama_sampler_k_shift *) smpl->ctx; auto * ctx = (llama_sampler_k_shift *) smpl->ctx;
if (ctx->k_set == true // ensures that k-shift can happen on the first step only
|| ctx->k <= 0 if (ctx->k_set != true) {
|| ctx->k >= (int) cur_p->size) { ctx->k_set = true;
} else {
return;
}
if (ctx->k <= 0 || ctx->k >= (int) cur_p->size) {
return; return;
} }
llama_sampler_top_shift_impl(cur_p, ctx->k); llama_sampler_top_shift_impl(cur_p, ctx->k);
ctx->k_set = true;
} }
static struct llama_sampler * llama_sampler_k_shift_clone(const struct llama_sampler * smpl) { static struct llama_sampler * llama_sampler_k_shift_clone(const struct llama_sampler * smpl) {