Fix a serious issue with addressing candidates by tokenid

Author: KerfuffleV2

llama.cpp

@@ -2825,27 +2825,32 @@ void llama_sample_seqrep_penalty(struct llama_context * ctx, llama_token_data_ar
     const bool ends_on_word = params->mid_word_scale == 1.0f
         || (llama_seqrep_check_word(ctx, last_tokens_p[last_tokens_size - 1]) & 2) != 0;
 
-    for (const auto it : penalize_tokens) {
-        const bool pt_starts_word = params->mid_word_scale == 1.0f ||
-            (llama_seqrep_check_word(ctx, it.first) & 1) != 0;
-        float scale = ends_on_word || pt_starts_word ? 1.0f : params->mid_word_scale;
+    for (size_t i = 0; i < candidates->size; ++i) {
+        auto pt_iter = penalize_tokens.find(candidates->data[i].id);
+        if (pt_iter == penalize_tokens.end()) {
+            continue;
+        }
 
-        float logit = candidates->data[it.first].logit;
+        const size_t count = pt_iter->second;
+        const bool pt_starts_word = params->mid_word_scale == 1.0f ||
+            (llama_seqrep_check_word(ctx, candidates->data[i].id) & 1) != 0;
+        float penalty_scale = ends_on_word || pt_starts_word ? 1.0f : params->mid_word_scale;
+        float logit = candidates->data[i].logit;
 
         if ((flags & LLAMA_SEQREP_DIVIDE_BY_PENALTY) == 0) {
             float penalty =
-                ( float(it.second) * params->length_penalty
-                + float(it.second > 0) * params->presence_penalty );
-            logit -= penalty * scale;
+                ( float(count) * params->length_penalty
+                + float(count > 0) * params->presence_penalty );
+            logit -= penalty * penalty_scale;
         } else {
             // This looks complicated. The point is to scale be able to scale penalties like
             // 1.2. For example, suppose length penalty is 1.2 and length is 3. 1.2 * 3 = 3.6
             // would be ridiculous. What we actually want is more like 1.6.
             // An alternative approach would be to iteratively apply the scale.
             // 10.0 / 1.6 == 6.25, however ((10.0 / 1.2) / 1.2) / 1.2 == 5.787
             float penalty =
-                ( (float(it.second) * (params->length_penalty - 1.0f))
-                + (float(it.second > 0) * (params->presence_penalty - 1.0f)) ) * scale
+                ( (float(count) * (params->length_penalty - 1.0f))
+                + (float(count > 0) * (params->presence_penalty - 1.0f)) ) * penalty_scale
                 + 1.0f;
             if (logit <= 0) {
                 logit *= penalty;
@@ -2857,7 +2862,7 @@ void llama_sample_seqrep_penalty(struct llama_context * ctx, llama_token_data_ar
                 logit = 0.0f;
             }
         }
-        candidates->data[it.first].logit = logit;
+        candidates->data[i].logit = logit;
     }
 
     candidates->sorted = false;