logit_bias: apply configurable escalating EOG bias at low n_remain

give eog an increasing (with length - per token, could be per
codepoint in future) bias, only after a configured amount generated

add to `sample_apply` an `n_remain` param, which is safer than having
logit_bias maintain state for how many times it's called (which
would lead to wrong assumptions e.g. when calling multiple times per
token).

see new command line options (incl a request 'after' instead of
'remain'):

-eog,  --eog-bias-per-tok N             when fewer than -start-eog-at-remain tokens are left to generate after
                                        -n, add this bias eog for each subsequent token (default: 0.0)
-remain, --start-eog-at-remain N        start applying -eog bias when this many tokens remain of the -n max
                                        (default: 0.0)
-after, --start-eog-after N             start applying -eog bias after this many tokens generated (default:
                                        1000000000.0); whichever happens first between -remain and -after
                                        applies

Verified that eog bias was effective at avoiding
overgeneration and is a reasonable supplement or alternative
to editing the prompt; a *constant* eog bias, already supported
in samplers, is likely to allow pathologically short outputs.

Author: graehl

common/arg.cpp

@@ -1656,6 +1656,15 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e
         exit(1); // for other exceptions, we exit with status code 1
     }
 
+    float &pafter = params.sampling.start_eog_after;
+    float &premain = params.sampling.start_eog_at_remain;
+    float const premain0 = premain;
+    float remain = params.n_predict - pafter;
+    if (premain < remain)
+      premain = remain;
+    if (params.sampling.eog_bias_per_tok)
+        LOG_INF("%s: n_predict=%d (first of start_eog_at_remain=%0.3g start_eog_after=%0.3g) => (remain=%0.3g) eog-bias-per-tok=%0.3g\n", __func__, (int) params.n_predict,
+                (double) premain0, (double) pafter, (double)premain, (double) params.sampling.eog_bias_per_tok);
     return true;
 }
 
@@ -2439,6 +2448,27 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             }
         }
     ).set_sparam());
+    add_opt(common_arg(
+        {"-eog", "--eog-bias-per-tok"}, "N",
+        string_format("when fewer than -start-eog-at-remain tokens are left to generate after -n, add this bias eog for each subsequent token (default: %.1f)", (double)params.sampling.eog_bias_per_tok),
+        [](common_params & params, const std::string & value) {
+            params.sampling.eog_bias_per_tok = std::stof(value);
+        }
+    ).set_sparam());
+    add_opt(common_arg(
+        {"-remain", "--start-eog-at-remain"}, "N",
+        string_format("start applying -eog bias when this many tokens remain of the -n max (default: %.1f)", (double)params.sampling.start_eog_at_remain),
+        [](common_params & params, const std::string & value) {
+            params.sampling.start_eog_at_remain = std::stof(value);
+        }
+    ).set_sparam());
+    add_opt(common_arg(
+        {"-after", "--start-eog-after"}, "N",
+        string_format("start applying -eog bias after this many tokens generated (default: %.1f); whichever happens first between -remain and -after applies", (double)params.sampling.start_eog_after),
+        [](common_params & params, const std::string & value) {
+            params.sampling.start_eog_after = std::stof(value);
+        }
+    ).set_sparam());
     add_opt(common_arg(
         {"--grammar"}, "GRAMMAR",
         string_format("BNF-like grammar to constrain generations (see samples in grammars/ dir) (default: '%s')", params.sampling.grammar.c_str()),

common/common.h

@@ -188,6 +188,13 @@ struct common_params_sampling {
     std::vector<llama_logit_bias> logit_bias;     // logit biases to apply
     std::vector<llama_logit_bias> logit_bias_eog; // pre-calculated logit biases for EOG tokens
 
+    float eog_bias_per_tok = 0; // escalating bias added to eog per token after:
+    /// this many remaining tokens (before applying eog_bias_per_tok) ...
+    float start_eog_at_remain = 0;
+    // or (whichever is first) after start_eog_after many generated:
+    /// (i.e. EOG logit bias = max(0,start_eog_after = max(start_eog_after, n_remain - start_eog_at_remain)) * eog_bias_per_tok)
+    float start_eog_after = 1e9;
+
     // print the parameters into a string
     std::string print() const;
 };

common/sampling.cpp

@@ -226,7 +226,10 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, co
             llama_sampler_init_logit_bias(
                 llama_vocab_n_tokens(vocab),
                 params.logit_bias.size(),
-                params.logit_bias.data()));
+                params.logit_bias.data(),
+                params.eog_bias_per_tok,
+                params.start_eog_at_remain,
+                vocab));
 
     if (params.mirostat == 0) {
         for (const auto & cnstr : params.samplers) {
@@ -336,18 +339,18 @@ void common_perf_print(const struct llama_context * ctx, const struct common_sam
     }
 }
 
-llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_context * ctx, int idx, bool grammar_first) {
+llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_context * ctx, int idx, bool grammar_first, float n_remain) {
     gsmpl->set_logits(ctx, idx);
 
     auto & grmr  = gsmpl->grmr;
     auto & chain = gsmpl->chain;
     auto & cur_p = gsmpl->cur_p; // initialized by set_logits
 
     if (grammar_first) {
-        llama_sampler_apply(grmr, &cur_p);
+        llama_sampler_apply(grmr, &cur_p, n_remain);
     }
 
-    llama_sampler_apply(chain, &cur_p);
+    llama_sampler_apply(chain, &cur_p, n_remain);
 
     GGML_ASSERT(cur_p.selected != -1 && "no selected token during sampling - check your sampling configuration");
 
@@ -362,7 +365,7 @@ llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_co
         llama_token_data       single_token_data       = { id, 1.0f, 0.0f };
         llama_token_data_array single_token_data_array = { &single_token_data, 1, -1, false };
 
-        llama_sampler_apply(grmr, &single_token_data_array);
+        llama_sampler_apply(grmr, &single_token_data_array, n_remain);
 
         const bool is_valid = single_token_data_array.data[0].logit != -INFINITY;
         if (is_valid) {
@@ -374,23 +377,23 @@ llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_co
     // if the token is not valid, sample again, but first apply the grammar sampler and then the sampling chain
     gsmpl->set_logits(ctx, idx);
 
-    llama_sampler_apply(grmr,  &cur_p);
-    llama_sampler_apply(chain, &cur_p);
+    llama_sampler_apply(grmr,  &cur_p, n_remain);
+    llama_sampler_apply(chain, &cur_p, n_remain);
 
     GGML_ASSERT(cur_p.selected != -1 && "no selected token during re-sampling - check your sampling configuration");
 
     return cur_p.data[cur_p.selected].id;
 }
 
-std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector<int> & idxs, const llama_tokens & draft, bool grammar_first) {
+std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector<int> & idxs, const llama_tokens & draft, bool grammar_first, float n_remain) {
     GGML_ASSERT(idxs.size() == draft.size() + 1 && "idxs.size() must be draft.size() + 1");
 
     std::vector<llama_token> result;
     result.reserve(idxs.size());
 
     size_t i = 0;
     for (; i < draft.size(); i++) {
-        const llama_token id = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first);
+        const llama_token id = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first, n_remain);
 
         common_sampler_accept(gsmpl, id, true);
 
@@ -402,7 +405,7 @@ std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sample
     }
 
     if (i == draft.size()) {
-        const llama_token id = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first);
+        const llama_token id = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first, n_remain);
 
         common_sampler_accept(gsmpl, id, true);
 
@@ -412,13 +415,13 @@ std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sample
     return result;
 }
 
-std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first) {
+std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first, float n_remain) {
     std::vector<int> idxs(draft.size() + 1);
     for (size_t i = 0; i < idxs.size(); ++i) {
         idxs[i] = i;
     }
 
-    return common_sampler_sample_and_accept_n(gsmpl, ctx, idxs, draft, grammar_first);
+    return common_sampler_sample_and_accept_n(gsmpl, ctx, idxs, draft, grammar_first, n_remain);
 }
 
 uint32_t common_sampler_get_seed(const struct common_sampler * gsmpl) {

common/sampling.h

@@ -58,7 +58,7 @@ void common_perf_print(const struct llama_context * ctx, const struct common_sam
 // if grammar_first is true, the grammar is applied before the samplers (slower)
 // useful in cases where all the resulting candidates (not just the sampled one) must fit the grammar
 //
-llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_context * ctx, int idx, bool grammar_first = false);
+llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_context * ctx, int idx, bool grammar_first = false, float n_remain = 0);
 
 // generalized version of common_sampler_sample
 //
@@ -76,10 +76,10 @@ llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_co
 //
 // returns at least 1 token, up to idxs.size()
 //
-std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector<int> & idxs, const llama_tokens & draft, bool grammar_first = false);
+std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector<int> & idxs, const llama_tokens & draft, bool grammar_first = false, float n_remain = 0);
 
 // assume idxs == [ 0, 1, 2, ..., draft.size() ]
-std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first = false);
+std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first = false, float n_remain = 0);
 
 uint32_t common_sampler_get_seed(const struct common_sampler * gsmpl);
 

common/speculative.cpp

@@ -310,12 +310,12 @@ llama_tokens common_speculative_gen_draft(
     llama_decode(ctx_dft, batch);
 
     common_sampler_reset(smpl);
-
+    int n_remain = params.n_draft;
     // sample n_draft tokens from the draft model
     for (int i = 0; i < params.n_draft; ++i) {
         common_batch_clear(batch);
 
-        common_sampler_sample(smpl, ctx_dft, 0, true);
+        common_sampler_sample(smpl, ctx_dft, 0, true, --n_remain);
 
         const auto * cur_p = common_sampler_get_candidates(smpl, true);
 

examples/batched/batched.cpp

@@ -162,7 +162,9 @@ int main(int argc, char ** argv) {
 
     const auto t_main_start = ggml_time_us();
 
+    int n_remain = n_predict;
     while (n_cur <= n_predict) {
+        --n_remain;
         // prepare the next batch
         common_batch_clear(batch);
 
@@ -173,7 +175,7 @@ int main(int argc, char ** argv) {
                 continue;
             }
 
-            const llama_token new_token_id = llama_sampler_sample(smpl, ctx, i_batch[i]);
+            const llama_token new_token_id = llama_sampler_sample(smpl, ctx, i_batch[i], n_remain);
 
             // is it an end of generation? -> mark the stream as finished
             if (llama_vocab_is_eog(vocab, new_token_id) || n_cur == n_predict) {

examples/diffusion/diffusion-cli.cpp

@@ -408,7 +408,8 @@ static void diffusion_generate(llama_context *          ctx,
                             false,
                         };
 
-                        llama_sampler_apply(sampler, &cur_p);
+                        float n_remain = params.max_length - pos;
+                        llama_sampler_apply(sampler, &cur_p, n_remain);
                         output_tokens[pos] = cur_p.data[cur_p.selected].id;
                     }
                 }
@@ -433,7 +434,8 @@ static void diffusion_generate(llama_context *          ctx,
                         false,
                     };
 
-                    llama_sampler_apply(sampler, &cur_p);
+                    float n_remain = params.max_length - i;
+                    llama_sampler_apply(sampler, &cur_p, n_remain);
                     llama_token sampled_token = cur_p.data[cur_p.selected].id;
 
                     float conf = calculate_confidence(cur_p, params.algorithm, rng);
@@ -477,7 +479,8 @@ static void diffusion_generate(llama_context *          ctx,
                         };
 
                         for (int32_t i = 0; i < std::min(transfer_count, (int32_t) confidences.size()); i++) {
-                            llama_sampler_apply(dist_sampler, &conf_array);
+                            float n_remain = params.max_length - i;
+                            llama_sampler_apply(dist_sampler, &conf_array, n_remain);
                             int32_t selected_idx = conf_array.selected;
                             int32_t mask_idx     = selected_idx;
                             int32_t pos          = mask_positions[mask_idx];

examples/lookahead/lookahead.cpp

@@ -253,6 +253,7 @@ int main(int argc, char ** argv) {
 
         int seq_id_best = 0;
 
+        int n_remain = N;
         for (int v = 0; v < N; ++v) {
             int i_batch = 0;
 
@@ -274,8 +275,9 @@ int main(int argc, char ** argv) {
                 }
             }
 
+            --n_remain;
             // sample the next token
-            id = common_sampler_sample(smpl, ctx, i_batch);
+            id = common_sampler_sample(smpl, ctx, i_batch, n_remain);
 
             common_sampler_accept(smpl, id, true);
 
@@ -349,10 +351,11 @@ int main(int argc, char ** argv) {
                     tokens_j[j] = tokens_j[j + 1];
                 }
 
+                unsigned constexpr NA = (unsigned)-1;
                 if (v == 0) {
                     // sample from the last level
                     for (int i = 0; i < W; i++) {
-                        tokens_j[N - 2][i] = common_sampler_sample(smpl, ctx, ngrams_cur.size()*(N-1) + W*(N - 2) + i);
+                      tokens_j[N - 2][i] = common_sampler_sample(smpl, ctx, ngrams_cur.size()*(N-1) + W*(N - 2) + i, NA);
                     }
                 } else {
                     for (int i = 0; i < W; i++) {

examples/lookup/lookup.cpp

@@ -117,7 +117,8 @@ int main(int argc, char ** argv){
         int i_dft = 0;
         while (true) {
             // sample from the target model
-            llama_token id = common_sampler_sample(smpl, ctx, i_dft);
+            unsigned const n_remain = params.n_predict - n_predict;
+            llama_token id = common_sampler_sample(smpl, ctx, i_dft, n_remain);
 
             common_sampler_accept(smpl, id, true);
 

examples/passkey/passkey.cpp

@@ -217,10 +217,12 @@ int main(int argc, char ** argv) {
 
     const auto t_main_start = ggml_time_us();
 
+    int n_remain = n_len - n_cur;
     while (n_cur <= n_len) {
+        --n_remain;
         // sample the next token
         {
-            const llama_token new_token_id = llama_sampler_sample(smpl, ctx, batch.n_tokens - 1);
+          const llama_token new_token_id = llama_sampler_sample(smpl, ctx, batch.n_tokens - 1, n_remain);
 
             // is it an end of generation?
             if (llama_vocab_is_eog(vocab, new_token_id) || n_cur == n_len) {