[Fix] Repetition penalty calculation fix (#2937)

Repetition penalty compilation fix.
mlc-ai · Sep 27, 2024 · fe79e9a · fe79e9a
1 parent 268d52a
commit fe79e9a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/mlc_llm/compiler_pass/attach_logit_processor.py b/python/mlc_llm/compiler_pass/attach_logit_processor.py
@@ -119,7 +119,7 @@ def _apply_penalty_inplace(  # pylint: disable=too-many-arguments,too-many-local
                         penalties[pos2seq_id[vp], 0] + token_cnt[vp] * penalties[pos2seq_id[vp], 1]
                     )
                     logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] = T.if_then_else(
-                        logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] > 0,
+                        logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] < 0,
                         logits[seq_ids[pos2seq_id[vp]], token_ids[vp]]
                         * penalties[pos2seq_id[vp], 2],
                         logits[seq_ids[pos2seq_id[vp]], token_ids[vp]]