bleam 0.0.9 → 0.0.11

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "bleam",
-  "version": "0.0.9",
+  "version": "0.0.11",
   "packageManager": "yarn@1.22.22",
   "type": "module",
   "files": [
@@ -39,6 +39,7 @@
     "react-native": "0.86.0",
     "tsx": "^4.21.0",
     "typescript": "~6.0.3",
+    "wrangler": "4.110.0",
     "ws": "^8.21.0",
     "yocto-spinner": "^1.2.1",
     "yoctocolors": "^2.1.2"

package/templates/image-generation/app/index.tsx

@@ -1,4 +1,4 @@
-import { generateImage } from 'bleam/ai'
+import { bonsai, generateImage } from 'bleam/ai'
 import { sx } from 'bleam/styles'
 import { Button } from 'bleam/ui'
 import { Window } from 'bleam/window'
@@ -6,6 +6,7 @@ import { useState } from 'react'
 import { Image, Text, View } from 'react-native'
 
 const appName = 'Bleam Basic'
+const model = bonsai('ternary-bonsai-image-4b')
 
 export default function App() {
   const [uri, setUri] = useState<string>()
@@ -15,6 +16,7 @@ export default function App() {
     setStatus('Generating...')
     try {
       const result = await generateImage({
+        model,
         prompt: 'A precise editorial illustration of a small native Mac app',
         aspectRatio: '1:1',
       })

package/templates/native/ios/Bleam/AI/Flux2/Configuration/ModelRegistry.swift

@@ -152,7 +152,7 @@ public enum ModelRegistry {
         case klein4B_bf16 = "klein4b-bf16"
         case klein4B_8bit = "klein4b-8bit"
         case klein4B_4bit = "klein4b-4bit"
-        case bonsaiKlein4B_2bit = "bonsai-klein4b-2bit"
+        case bonsaiKlein4B_2bit = "ternary-bonsai-image-4b"
         case bonsaiKlein4B_1bit = "bonsai-klein4b-1bit"
         case klein9B_8bit = "klein9b-8bit"
         case klein9B_4bit = "klein9b-4bit"

package/templates/native/ios/Bleam/AI/Flux2/Pipeline/FluxImageRunner.swift

@@ -198,12 +198,17 @@ public final class FluxImageRunner: @unchecked Sendable {
     private static let referenceImageMaxLongSide = 768
 
     private let paths: FluxImageRunnerPaths
+    private let transformerProfile: Flux2TransformerRepoProfile
     private let scheduler = FlowMatchEulerScheduler()
     private var transformer: Flux2Transformer2DModel?
     private var vae: AutoencoderKLFlux2?
 
-    public init(paths: FluxImageRunnerPaths) {
+    public init(
+        paths: FluxImageRunnerPaths,
+        transformerProfile: Flux2TransformerRepoProfile = ModelRegistry.TransformerRepo.moxinOrgKlein4B8bit.profile
+    ) {
         self.paths = paths
+        self.transformerProfile = transformerProfile
     }
 
     public func generate(
@@ -1237,7 +1242,7 @@ public final class FluxImageRunner: @unchecked Sendable {
             return transformer
         }
 
-        let profile = ModelRegistry.TransformerRepo.moxinOrgKlein4B8bit.profile
+        let profile = transformerProfile
         let configURL = paths.transformer.appendingPathComponent("config.json")
         let config = FileManager.default.fileExists(atPath: configURL.path)
             ? (try? Flux2TransformerConfig.load(from: configURL)) ?? .klein4B

package/templates/native/ios/Bleam/AI/TextGeneration/Bonsai/BonsaiTextGenerator.swift

@@ -0,0 +1,46 @@
+import Foundation
+import MLX
+
+public enum BonsaiTextGenerator {
+    public static func generate(
+        profile: BonsaiTextModelProfile,
+        modelURL: URL,
+        inputIDs: [Int],
+        endTokenIDs: Set<Int>? = nil,
+        parameters: GenerationParameters? = nil,
+        isCancelled: @escaping @Sendable () -> Bool,
+        onToken: @escaping @Sendable (Int) -> Bool
+    ) async throws -> TokenGenerationResult {
+        try BonsaiTextSnapshotValidator.validate(modelURL: modelURL, profile: profile)
+        let parameters = parameters ?? profile.defaultGenerationParameters
+        let endTokenIDs = endTokenIDs ?? profile.defaultEndTokenIDs
+        defer { MLX.Memory.clearCache() }
+
+        switch profile.backend {
+        case .qwen3:
+            try Task.checkCancellation()
+            let model = try Qwen3CheckpointLoader.load(from: modelURL)
+            try Task.checkCancellation()
+            return try Qwen3InferenceSession(model: model, endTokenIDs: endTokenIDs).generate(
+                inputIDs: inputIDs,
+                parameters: parameters,
+                isCancelled: { Task.isCancelled || isCancelled() },
+                onToken: onToken
+            )
+        case .qwen35:
+            let model = try Qwen35Checkpoint.load(
+                from: modelURL,
+                expectedQuantization: Qwen35Quantization(groupSize: 128, bits: 2)
+            )
+            let session = Qwen35Session(model: model)
+            defer { session.reset() }
+            return try await session.generate(
+                inputIDs: inputIDs,
+                parameters: parameters,
+                endTokenIDs: endTokenIDs,
+                isCancelled: { Task.isCancelled || isCancelled() },
+                onToken: onToken
+            )
+        }
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Bonsai/BonsaiTextModel.swift

@@ -0,0 +1,107 @@
+import Foundation
+
+public enum BonsaiTextModelID: String, CaseIterable, Codable, Sendable {
+    case ternaryBonsai17B = "ternary-bonsai-1.7b"
+    case ternaryBonsai4B = "ternary-bonsai-4b"
+    case ternaryBonsai8B = "ternary-bonsai-8b"
+    case ternaryBonsai27B = "ternary-bonsai-27b"
+}
+
+public enum BonsaiTextBackend: Sendable {
+    case qwen3
+    case qwen35
+}
+
+public struct BonsaiTextModelProfile: Sendable {
+    public let id: BonsaiTextModelID
+    public let backend: BonsaiTextBackend
+    public let vocabularySize: Int
+    public let hiddenSize: Int
+    public let intermediateSize: Int
+    public let hiddenLayerCount: Int
+    public let attentionHeadCount: Int
+    public let keyValueHeadCount: Int
+    public let headDimension: Int
+    public let contextLength: Int
+    public let defaultEndTokenIDs: Set<Int>
+    public let defaultGenerationParameters: GenerationParameters
+
+    public static func resolve(_ value: String) throws -> Self {
+        guard let id = BonsaiTextModelID(rawValue: value) else {
+            throw BonsaiTextError.unsupportedModel(value)
+        }
+        return profile(for: id)
+    }
+
+    public static func profile(for id: BonsaiTextModelID) -> Self {
+        switch id {
+        case .ternaryBonsai17B:
+            return Self(
+                id: id, backend: .qwen3, vocabularySize: 151_669, hiddenSize: 2_048,
+                intermediateSize: 6_144, hiddenLayerCount: 28, attentionHeadCount: 16,
+                keyValueHeadCount: 8, headDimension: 128, contextLength: 32_768,
+                defaultEndTokenIDs: [151_643, 151_645],
+                defaultGenerationParameters: GenerationParameters(
+                    maximumNewTokens: 256, temperature: 0.5, topP: 0.85,
+                    repetitionPenalty: 1, repetitionContextSize: 64
+                )
+            )
+        case .ternaryBonsai4B:
+            return Self(
+                id: id, backend: .qwen3, vocabularySize: 151_669, hiddenSize: 2_560,
+                intermediateSize: 9_728, hiddenLayerCount: 36, attentionHeadCount: 32,
+                keyValueHeadCount: 8, headDimension: 128, contextLength: 32_768,
+                defaultEndTokenIDs: [151_643, 151_645],
+                defaultGenerationParameters: GenerationParameters(
+                    maximumNewTokens: 256, temperature: 0.5, topP: 0.85,
+                    repetitionPenalty: 1, repetitionContextSize: 64
+                )
+            )
+        case .ternaryBonsai8B:
+            return Self(
+                id: id, backend: .qwen3, vocabularySize: 151_669, hiddenSize: 4_096,
+                intermediateSize: 12_288, hiddenLayerCount: 36, attentionHeadCount: 32,
+                keyValueHeadCount: 8, headDimension: 128, contextLength: 65_536,
+                defaultEndTokenIDs: [151_643, 151_645],
+                defaultGenerationParameters: GenerationParameters(
+                    maximumNewTokens: 256, temperature: 0.5, topP: 0.85,
+                    repetitionPenalty: 1, repetitionContextSize: 64
+                )
+            )
+        case .ternaryBonsai27B:
+            return Self(
+                id: id, backend: .qwen35, vocabularySize: 248_320, hiddenSize: 5_120,
+                intermediateSize: 17_408, hiddenLayerCount: 64, attentionHeadCount: 24,
+                keyValueHeadCount: 4, headDimension: 256, contextLength: 262_144,
+                defaultEndTokenIDs: [248_044, 248_046],
+                defaultGenerationParameters: GenerationParameters(
+                    maximumNewTokens: 2_048, temperature: 0.7, topP: 0.95, topK: 20,
+                    repetitionPenalty: 1, repetitionContextSize: 128
+                )
+            )
+        }
+    }
+}
+
+public enum BonsaiTextError: LocalizedError {
+    case unsupportedModel(String)
+    case invalidModelPath(String)
+    case snapshotMismatch(model: BonsaiTextModelID, detail: String)
+    case invalidSeed(Int)
+    case invalidGenerationParameters(String)
+
+    public var errorDescription: String? {
+        switch self {
+        case .unsupportedModel(let value):
+            "Unsupported Bonsai text model: \(value)"
+        case .invalidModelPath(let path):
+            "Bonsai model path is not a directory: \(path)"
+        case .snapshotMismatch(let model, let detail):
+            "Snapshot does not match \(model.rawValue): \(detail)"
+        case .invalidSeed(let seed):
+            "Seed must be nonnegative, got \(seed)"
+        case .invalidGenerationParameters(let detail):
+            "Invalid generation parameters: \(detail)"
+        }
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Bonsai/BonsaiTextSnapshotValidator.swift

@@ -0,0 +1,201 @@
+import Foundation
+
+public enum BonsaiTextSnapshotValidator {
+    private struct Quantization: Decodable, Equatable {
+        let groupSize: Int
+        let bits: Int
+        let mode: String
+
+        private enum CodingKeys: String, CodingKey {
+            case groupSize = "group_size"
+            case bits
+            case mode
+        }
+
+        init(groupSize: Int, bits: Int, mode: String = "affine") {
+            self.groupSize = groupSize
+            self.bits = bits
+            self.mode = mode
+        }
+
+        init(from decoder: Decoder) throws {
+            let values = try decoder.container(keyedBy: CodingKeys.self)
+            groupSize = try values.decode(Int.self, forKey: .groupSize)
+            bits = try values.decode(Int.self, forKey: .bits)
+            mode = try values.decodeIfPresent(String.self, forKey: .mode) ?? "affine"
+        }
+    }
+
+    private struct Qwen3Identity: Decodable {
+        struct RopeScaling: Decodable {
+            let ropeType: String
+            let factor: Float
+            let originalMaximumPositions: Int
+
+            private enum CodingKeys: String, CodingKey {
+                case ropeType = "rope_type"
+                case factor
+                case originalMaximumPositions = "original_max_position_embeddings"
+            }
+        }
+
+        let hiddenActivation: String
+        let layerTypes: [String]
+        let ropeScaling: RopeScaling
+        let quantization: Quantization
+
+        private enum CodingKeys: String, CodingKey {
+            case hiddenActivation = "hidden_act"
+            case layerTypes = "layer_types"
+            case ropeScaling = "rope_scaling"
+            case quantization
+        }
+    }
+
+    private struct Qwen35TextIdentity: Decodable {
+        let modelType: String
+        let vocabularySize: Int
+        let hiddenSize: Int
+        let intermediateSize: Int
+        let hiddenLayerCount: Int
+        let attentionHeadCount: Int
+        let keyValueHeadCount: Int
+        let headDimension: Int
+        let maximumPositionCount: Int
+        let quantization: Quantization
+
+        private enum CodingKeys: String, CodingKey {
+            case modelType = "model_type"
+            case vocabularySize = "vocab_size"
+            case hiddenSize = "hidden_size"
+            case intermediateSize = "intermediate_size"
+            case hiddenLayerCount = "num_hidden_layers"
+            case attentionHeadCount = "num_attention_heads"
+            case keyValueHeadCount = "num_key_value_heads"
+            case headDimension = "head_dim"
+            case maximumPositionCount = "max_position_embeddings"
+            case quantization
+        }
+    }
+
+    private struct Qwen35Identity: Decodable {
+        let architectures: [String]
+        let modelType: String
+        let textConfiguration: Qwen35TextIdentity
+        let quantization: Quantization
+
+        private enum CodingKeys: String, CodingKey {
+            case architectures
+            case modelType = "model_type"
+            case textConfiguration = "text_config"
+            case quantization
+        }
+    }
+
+    public static func validate(modelURL: URL, profile: BonsaiTextModelProfile) throws {
+        var isDirectory: ObjCBool = false
+        guard FileManager.default.fileExists(atPath: modelURL.path, isDirectory: &isDirectory),
+              isDirectory.boolValue else {
+            throw BonsaiTextError.invalidModelPath(modelURL.path)
+        }
+
+        do {
+            switch profile.backend {
+            case .qwen3:
+                try validateQwen3(modelURL: modelURL, profile: profile)
+            case .qwen35:
+                try validateQwen35(modelURL: modelURL, profile: profile)
+            }
+        } catch let error as BonsaiTextError {
+            throw error
+        } catch {
+            throw BonsaiTextError.snapshotMismatch(model: profile.id, detail: error.localizedDescription)
+        }
+    }
+
+    private static func validateQwen3(modelURL: URL, profile: BonsaiTextModelProfile) throws {
+        let configuration = try Qwen3ModelConfiguration.load(from: modelURL)
+        try configuration.validate()
+        let data = try Data(contentsOf: modelURL.appendingPathComponent("config.json"))
+        let identity = try JSONDecoder().decode(Qwen3Identity.self, from: data)
+        guard configuration.architectures == ["Qwen3ForCausalLM"], configuration.modelType == "qwen3" else {
+            throw mismatch(profile, "expected Qwen3ForCausalLM/qwen3 architecture")
+        }
+        guard configuration.vocabularySize == profile.vocabularySize,
+              configuration.hiddenSize == profile.hiddenSize,
+              configuration.intermediateSize == profile.intermediateSize,
+              configuration.hiddenLayerCount == profile.hiddenLayerCount,
+              configuration.attentionHeadCount == profile.attentionHeadCount,
+              configuration.keyValueHeadCount == profile.keyValueHeadCount,
+              configuration.headDimension == profile.headDimension,
+              configuration.maximumPositionCount == profile.contextLength else {
+            throw mismatch(profile, "Qwen3 dimensions do not match the stable model ID")
+        }
+        let expectedRopeTheta: Float = switch profile.id {
+        case .ternaryBonsai4B: 5_000_000
+        default: 1_000_000
+        }
+        guard configuration.rmsNormEpsilon == 1e-6,
+              configuration.ropeTheta == expectedRopeTheta,
+              configuration.attentionBias == false,
+              configuration.tieWordEmbeddings == (profile.id != .ternaryBonsai8B),
+              identity.hiddenActivation == "silu",
+              identity.layerTypes.count == profile.hiddenLayerCount,
+              identity.layerTypes.allSatisfy({ $0 == "full_attention" }),
+              identity.ropeScaling.ropeType == "yarn",
+              identity.ropeScaling.factor == 4,
+              identity.ropeScaling.originalMaximumPositions == profile.contextLength / 4 else {
+            throw mismatch(profile, "Qwen3 architecture settings do not match the stable model ID")
+        }
+        guard configuration.quantization?.bits == 2,
+              configuration.quantization?.groupSize == 128,
+              identity.quantization == Quantization(groupSize: 128, bits: 2) else {
+            throw mismatch(profile, "expected packed 2-bit/group-128 quantization")
+        }
+    }
+
+    private static func validateQwen35(modelURL: URL, profile: BonsaiTextModelProfile) throws {
+        let data = try Data(contentsOf: modelURL.appendingPathComponent("config.json"))
+        let identity = try JSONDecoder().decode(Qwen35Identity.self, from: data)
+        let text = identity.textConfiguration
+        guard identity.architectures == ["Qwen3_5ForConditionalGeneration"],
+              identity.modelType == "qwen3_5", text.modelType == "qwen3_5_text" else {
+            throw mismatch(profile, "expected Qwen3.5 conditional-generation/text architecture")
+        }
+        guard text.vocabularySize == profile.vocabularySize,
+              text.hiddenSize == profile.hiddenSize,
+              text.intermediateSize == profile.intermediateSize,
+              text.hiddenLayerCount == profile.hiddenLayerCount,
+              text.attentionHeadCount == profile.attentionHeadCount,
+              text.keyValueHeadCount == profile.keyValueHeadCount,
+              text.headDimension == profile.headDimension,
+              text.maximumPositionCount == profile.contextLength else {
+            throw mismatch(profile, "Qwen3.5 dimensions do not match the stable model ID")
+        }
+        let requiredQuantization = Quantization(groupSize: 128, bits: 2)
+        guard identity.quantization == requiredQuantization,
+              text.quantization == requiredQuantization else {
+            throw mismatch(profile, "Qwen3.5 requires 2-bit/group-128 at both snapshot levels")
+        }
+
+        let backendConfiguration = try Qwen35Configuration.load(from: modelURL)
+        try backendConfiguration.validate(expectedQuantization: Qwen35Quantization(groupSize: 128, bits: 2))
+        guard backendConfiguration.rmsNormEps == 1e-6,
+              backendConfiguration.attentionBias == false,
+              backendConfiguration.tieWordEmbeddings == false,
+              backendConfiguration.fullAttentionInterval == 4,
+              backendConfiguration.linearConvKernelDim == 4,
+              backendConfiguration.linearKeyHeadDim == 128,
+              backendConfiguration.linearValueHeadDim == 128,
+              backendConfiguration.linearNumKeyHeads == 16,
+              backendConfiguration.linearNumValueHeads == 48,
+              backendConfiguration.ropeTheta == 10_000_000,
+              backendConfiguration.partialRotaryFactor == 0.25 else {
+            throw mismatch(profile, "Qwen3.5 architecture settings do not match the stable model ID")
+        }
+    }
+
+    private static func mismatch(_ profile: BonsaiTextModelProfile, _ detail: String) -> BonsaiTextError {
+        .snapshotMismatch(model: profile.id, detail: detail)
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Core/GenerationTypes.swift

@@ -0,0 +1,85 @@
+import Foundation
+
+public struct GenerationParameters: Sendable {
+    public var maximumNewTokens: Int
+    public var temperature: Float
+    public var topP: Float
+    public var topK: Int?
+    public var repetitionPenalty: Float
+    public var repetitionContextSize: Int
+    public var seed: UInt64?
+
+    public init(
+        maximumNewTokens: Int = 512,
+        temperature: Float = 0.7,
+        topP: Float = 0.95,
+        topK: Int? = nil,
+        repetitionPenalty: Float = 1,
+        repetitionContextSize: Int = 64,
+        seed: UInt64? = nil
+    ) {
+        self.maximumNewTokens = maximumNewTokens
+        self.temperature = temperature
+        self.topP = topP
+        self.topK = topK
+        self.repetitionPenalty = repetitionPenalty
+        self.repetitionContextSize = repetitionContextSize
+        self.seed = seed
+    }
+
+    public static let balanced = GenerationParameters()
+}
+
+public enum GenerationFinishReason: String, Codable, Sendable {
+    case endOfSequence
+    case length
+    case cancelled
+    case callbackStopped
+}
+
+public struct TokenGenerationResult: Sendable {
+    public let promptTokenCount: Int
+    public let tokenIDs: [Int]
+    public let finishReason: GenerationFinishReason
+
+    public init(promptTokenCount: Int, tokenIDs: [Int], finishReason: GenerationFinishReason) {
+        self.promptTokenCount = promptTokenCount
+        self.tokenIDs = tokenIDs
+        self.finishReason = finishReason
+    }
+}
+
+public enum TextGenerationError: LocalizedError {
+    case emptyPrompt
+    case invalidMaximumNewTokens(Int)
+    case contextExceeded(requested: Int, maximum: Int)
+    case invalidTokenID(Int)
+    case invalidTemperature(Float)
+    case invalidTopP(Float)
+    case invalidTopK(Int)
+    case invalidRepetitionPenalty(Float)
+    case invalidRepetitionContextSize(Int)
+
+    public var errorDescription: String? {
+        switch self {
+        case .emptyPrompt:
+            "The pre-tokenized prompt is empty"
+        case .invalidMaximumNewTokens(let value):
+            "Maximum new tokens must be nonnegative, got \(value)"
+        case .contextExceeded(let requested, let maximum):
+            "Requested context length \(requested) exceeds model maximum \(maximum)"
+        case .invalidTokenID(let value):
+            "Token ID \(value) is outside the model vocabulary"
+        case .invalidTemperature(let value):
+            "Temperature must be nonnegative, got \(value)"
+        case .invalidTopP(let value):
+            "Top-p must be in (0, 1], got \(value)"
+        case .invalidTopK(let value):
+            "Top-k must be positive, got \(value)"
+        case .invalidRepetitionPenalty(let value):
+            "Repetition penalty must be positive, got \(value)"
+        case .invalidRepetitionContextSize(let value):
+            "Repetition context size must be nonnegative, got \(value)"
+        }
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Core/ModelCapabilities.swift

@@ -0,0 +1,21 @@
+import MLX
+
+/// Encoding is an independent capability; callers do not need a causal LM head.
+public protocol HiddenStateModel: AnyObject {
+    var vocabularySize: Int { get }
+    var maximumContextLength: Int { get }
+
+    func hiddenStates(
+        for inputIDs: MLXArray,
+        at layerIndices: Set<Int>,
+        attentionMask: MLXArray?
+    ) -> [Int: MLXArray]
+}
+
+public protocol CausalLanguageModel: AnyObject {
+    var vocabularySize: Int { get }
+    var maximumContextLength: Int { get }
+
+    func makeCache() -> [TextKVCache]
+    func logits(for inputIDs: MLXArray, cache: [TextKVCache]) -> MLXArray
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Core/TextKVCache.swift

@@ -0,0 +1,23 @@
+import MLX
+
+public final class TextKVCache {
+    private var keys: MLXArray?
+    private var values: MLXArray?
+
+    public private(set) var length = 0
+
+    public init() {}
+
+    public func update(keys newKeys: MLXArray, values newValues: MLXArray) -> (MLXArray, MLXArray) {
+        keys = keys.map { concatenated([$0, newKeys], axis: 2) } ?? newKeys
+        values = values.map { concatenated([$0, newValues], axis: 2) } ?? newValues
+        length += newKeys.dim(2)
+        return (keys!, values!)
+    }
+
+    public func clear() {
+        keys = nil
+        values = nil
+        length = 0
+    }
+}

package/templates/native/ios/Bleam/AI/TextGeneration/Core/TokenSampler.swift

@@ -0,0 +1,79 @@
+import MLX
+
+public enum TokenSampler {
+    public static func validate(_ parameters: GenerationParameters) throws {
+        guard parameters.maximumNewTokens >= 0 else {
+            throw TextGenerationError.invalidMaximumNewTokens(parameters.maximumNewTokens)
+        }
+        guard parameters.temperature >= 0 else {
+            throw TextGenerationError.invalidTemperature(parameters.temperature)
+        }
+        guard parameters.topP > 0, parameters.topP <= 1 else {
+            throw TextGenerationError.invalidTopP(parameters.topP)
+        }
+        if let topK = parameters.topK, topK <= 0 {
+            throw TextGenerationError.invalidTopK(topK)
+        }
+        guard parameters.repetitionPenalty > 0 else {
+            throw TextGenerationError.invalidRepetitionPenalty(parameters.repetitionPenalty)
+        }
+        guard parameters.repetitionContextSize >= 0 else {
+            throw TextGenerationError.invalidRepetitionContextSize(parameters.repetitionContextSize)
+        }
+    }
+
+    public static func sample(
+        logits: MLXArray,
+        parameters: GenerationParameters,
+        tokenHistory: [Int]
+    ) throws -> MLXArray {
+        try validate(parameters)
+        let selected: MLXArray
+        switch logits.ndim {
+        case 3: selected = logits[0, -1]
+        case 2: selected = logits[0]
+        default: selected = logits
+        }
+        let values = selected
+
+        if parameters.repetitionPenalty != 1,
+           parameters.repetitionContextSize > 0,
+           !tokenHistory.isEmpty {
+            let recent = Set(tokenHistory.suffix(parameters.repetitionContextSize))
+            let valid = recent.filter { $0 >= 0 && $0 < values.dim(-1) }
+            if !valid.isEmpty {
+                let indices = MLXArray(valid.sorted().map(Int32.init))
+                let repeated = MLX.take(values, indices, axis: -1)
+                values[indices] = MLX.where(
+                    repeated .> 0,
+                    repeated / parameters.repetitionPenalty,
+                    repeated * parameters.repetitionPenalty
+                )
+            }
+        }
+
+        if parameters.temperature == 0 {
+            return argMax(values).asType(Int32.self)
+        }
+
+        // Match mlx-lm ordering: select top-p from untempered probabilities.
+        let probabilities = softmax(values, axis: -1)
+        let candidateIndices: MLXArray
+        if let topK = parameters.topK, topK < probabilities.dim(-1) {
+            candidateIndices = argPartition(probabilities, kth: -topK, axis: -1)[(-topK)...]
+        } else {
+            candidateIndices = MLXArray.arange(probabilities.dim(-1))
+        }
+        let candidates = take(probabilities, candidateIndices, axis: -1)
+        let order = argSort(-candidates, axis: -1)
+        let sortedIndices = take(candidateIndices, order, axis: -1)
+        let sortedProbabilities = take(candidates, order, axis: -1)
+        let cumulativeBefore = cumsum(sortedProbabilities, axis: -1) - sortedProbabilities
+        let filtered = MLX.where(
+            cumulativeBefore .< parameters.topP,
+            log(sortedProbabilities) / parameters.temperature,
+            MLX.full(sortedProbabilities.shape, values: -Float.infinity)
+        )
+        return sortedIndices[MLXRandom.categorical(filtered)].asType(Int32.self)
+    }
+}