chore: add repetition_penalty example (#45)

Libraries/LLM/Evaluate.swift

@@ -25,6 +25,25 @@ private func topPSampling(logits: MLXArray, topP: Float, temp: Float) -> MLXArra
     return sortedIndices.squeezed(axis: 0)[sortedToken]
 }
 
+private func applyRepetitionPenalty(
+    logits: MLXArray, repetitionContext: MLXArray, penalty: Float
+) -> MLXArray {
+    var logits = logits
+
+    if repetitionContext.shape[0] > 0 {
+        let indices = repetitionContext
+        var selectedLogits = take(logits, indices, axis: -1).squeezed(axis: 0)
+
+        selectedLogits = MLX.where(
+            selectedLogits .< 0, selectedLogits * penalty, selectedLogits / penalty)
+
+        logits[0..., indices] = selectedLogits
+        return logits
+    }
+
+    return logits
+}
+
 private func sample(logits: MLXArray, temp: Float, topP: Float = 1.0) -> MLXArray {
     if temp == 0 {
         return argMax(logits, axis: -1)
@@ -43,23 +62,53 @@ public struct TokenIterator: Sequence, IteratorProtocol {
     let model: LLMModel
     let temp: Float
     let topP: Float
+    let repetitionPenalty: Float
+    let repetitionContextSize: Int
+    var repetitionContext: MLXArray
     var y: MLXArray
     var cache: [(MLXArray, MLXArray)]
 
     var first = true
 
-    public init(prompt: MLXArray, model: LLMModel, temp: Float = 0.0, topP: Float = 1.0) {
+    public init(
+        prompt: MLXArray, model: LLMModel, temp: Float = 0.0, topP: Float = 1.0,
+        repetitionPenalty: Float = 1.0, repetitionContextSize: Int = 20
+    ) {
         self.model = model
         self.temp = temp
         self.topP = topP
         self.y = prompt
         self.cache = []
+        self.repetitionPenalty = repetitionPenalty
+        self.repetitionContextSize = repetitionContextSize
+        if repetitionContextSize > 1 {
+            if prompt.shape[0] <= repetitionContextSize {
+                self.repetitionContext = prompt
+            } else {
+                self.repetitionContext = prompt[-repetitionContextSize ... -1]
+            }
+        } else {
+            self.repetitionContext = []
+        }
     }
 
     mutating public func next() -> MLXArray? {
         var logits: MLXArray
         (logits, cache) = model(expandedDimensions(y, axis: 0), cache: cache.isEmpty ? nil : cache)
-        y = sample(logits: logits[-1, axis: 1], temp: temp, topP: topP)
+        logits = logits[0..., -1, 0...]
+        if repetitionPenalty > 1.0 {
+            // apply repetition penalty
+            logits = applyRepetitionPenalty(
+                logits: logits, repetitionContext: repetitionContext, penalty: repetitionPenalty)
+        }
+        y = sample(logits: logits, temp: temp, topP: topP)
+        // append the current token to the context and check repetitionPenalty context see if need to remove the first token
+        if repetitionContextSize > 1 {
+            repetitionContext = concatenated([repetitionContext, y], axis: 0)
+            if repetitionContext.shape[0] > repetitionContextSize {
+                repetitionContext = repetitionContext[1...]
+            }
+        }
 
         return y
     }
@@ -71,7 +120,10 @@ public struct TokenIterator: Sequence, IteratorProtocol {
 ///
 /// Note that because MLXArray is not thread safe this eval's the result and sends the TokenId back
 /// to the caller.
-public func generate(prompt: MLXArray, model: LLMModel, temp: Float = 0.0, topP: Float = 1.0) -> (
+public func generate(
+    prompt: MLXArray, model: LLMModel, temp: Float = 0.0, topP: Float = 1.0,
+    repetitionPenalty: Float = 1.0, repetitionContextSize: Int = 20
+) -> (
     Task<Void, Never>, AsyncBufferSequence<AsyncChannel<Int>>
 ) {
     let channel = AsyncChannel<Int>()
@@ -80,12 +132,38 @@ public func generate(prompt: MLXArray, model: LLMModel, temp: Float = 0.0, topP:
     let task = Task {
         var y = prompt
         var cache = [(MLXArray, MLXArray)]()
+        var repetitionContext: MLXArray
 
+        if repetitionContextSize > 1 {
+            if prompt.shape[0] <= repetitionContextSize {
+                repetitionContext = prompt
+            } else {
+                repetitionContext = prompt[-repetitionContextSize ... -1]
+            }
+        } else {
+            repetitionContext = []
+        }
         while !Task.isCancelled {
             var logits: MLXArray
             (logits, cache) = model(
                 expandedDimensions(y, axis: 0), cache: cache.isEmpty ? nil : cache)
-            y = sample(logits: logits[-1, axis: 1], temp: temp, topP: topP)
+
+            logits = logits[0..., -1, 0...]
+            if repetitionPenalty > 1.0 {
+                // apply repetition penalty
+                logits = applyRepetitionPenalty(
+                    logits: logits, repetitionContext: repetitionContext, penalty: repetitionPenalty
+                )
+            }
+            y = sample(logits: logits, temp: temp, topP: topP)
+            // append the current token to the context and check repetitionPenalty context see if need to remove the first token
+            if repetitionContextSize > 1 {
+                repetitionContext = concatenated([repetitionContext, y], axis: 0)
+                if repetitionContext.shape[0] > repetitionContextSize {
+                    repetitionContext = repetitionContext[1...]
+                }
+            }
+
             eval(y)
 
             await channel.send(y.item(Int.self))