fix: implement repetition_penalty as logits processor — make_sampler lacks it

mlx-lm 0.31.2's make_sampler() does not accept repetition_penalty
(only temp, top_p, min_p, top_k, xtc_*). The previous fix passed it
as a keyword arg, breaking all generation with TypeError.

Now implemented as a proper logits processor that tracks recently
generated tokens and penalises their logits (divides positive logits,
multiplies negative logits by the penalty factor). Uses a sliding
context window of 100 tokens.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: longevityboris

scripts/mlx_generate.py

@@ -72,6 +72,28 @@ def processor(tokens, logits):
     return processor
 
 
+def _make_repetition_processor(penalty, context_size=100):
+    """Build a logits processor that penalises recently generated tokens."""
+    import mlx.core as mx
+
+    generated_tokens = []
+
+    def processor(tokens, logits):
+        generated_tokens.append(int(tokens[-1]) if tokens.size > 0 else 0)
+        recent = generated_tokens[-context_size:]
+        if recent:
+            ids = mx.array(list(set(recent)), dtype=mx.int32)
+            penalties = mx.where(
+                logits[..., ids] > 0,
+                logits[..., ids] / penalty,
+                logits[..., ids] * penalty,
+            )
+            logits[..., ids] = penalties
+        return logits
+
+    return processor
+
+
 def main():
     req = json.loads(sys.stdin.read())
 
@@ -118,24 +140,28 @@ def main():
                 formatted += system_prompt + "\n\n"
             formatted += prompt_text
 
-    sampler = make_sampler(
-        temp=temperature,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty if repetition_penalty > 1.0 else None,
-        repetition_context_size=100,
-    )
+    sampler = make_sampler(temp=temperature, top_p=top_p)
 
     # Build logits processor for vocabulary bias
     logits_processor = build_logits_processor(logit_bias, tokenizer)
 
+    # Chain repetition penalty as a logits processor
+    processors = []
+    if repetition_penalty > 1.0:
+        processors.append(
+            _make_repetition_processor(repetition_penalty, context_size=100)
+        )
+    if logits_processor is not None:
+        processors.append(logits_processor)
+
     t0 = time.time()
 
     gen_kwargs = dict(
         max_tokens=max_tokens,
         sampler=sampler,
     )
-    if logits_processor is not None:
-        gen_kwargs["logits_processors"] = [logits_processor]
+    if processors:
+        gen_kwargs["logits_processors"] = processors
 
     full_text = ""
     last_resp = None

scripts/mlx_worker.py

@@ -58,6 +58,28 @@ def processor(tokens, logits):
     return processor
 
 
+def _make_repetition_processor(penalty, context_size=100):
+    """Build a logits processor that penalises recently generated tokens."""
+    import mlx.core as mx
+
+    generated_tokens = []
+
+    def processor(tokens, logits):
+        generated_tokens.append(int(tokens[-1]) if tokens.size > 0 else 0)
+        recent = generated_tokens[-context_size:]
+        if recent:
+            ids = mx.array(list(set(recent)), dtype=mx.int32)
+            penalties = mx.where(
+                logits[..., ids] > 0,
+                logits[..., ids] / penalty,
+                logits[..., ids] * penalty,
+            )
+            logits[..., ids] = penalties
+        return logits
+
+    return processor
+
+
 def main():
     import mlx.core as mx
     from mlx_lm import load, stream_generate
@@ -137,17 +159,21 @@ def main():
                         formatted += system_prompt + "\n\n"
                     formatted += prompt_text
 
-            sampler = make_sampler(
-                temp=temperature,
-                top_p=top_p,
-                repetition_penalty=repetition_penalty if repetition_penalty > 1.0 else None,
-                repetition_context_size=100,
-            )
+            sampler = make_sampler(temp=temperature, top_p=top_p)
             logits_processor = build_logits_processor(logit_bias, tokenizer)
 
-            gen_kwargs = dict(max_tokens=max_tokens, sampler=sampler)
+            # Chain repetition penalty as a logits processor
+            processors = []
+            if repetition_penalty > 1.0:
+                processors.append(
+                    _make_repetition_processor(repetition_penalty, context_size=100)
+                )
             if logits_processor is not None:
-                gen_kwargs["logits_processors"] = [logits_processor]
+                processors.append(logits_processor)
+
+            gen_kwargs = dict(max_tokens=max_tokens, sampler=sampler)
+            if processors:
+                gen_kwargs["logits_processors"] = processors
 
             t0 = time.time()