Fix constrained sampling ignoring the allowed-token set

Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift

@@ -673,8 +673,7 @@ import Foundation
             params.n_gpu_layers = 0
 
             // Try to reduce memory usage
-            params.use_mmap = true
-            params.use_mlock = false
+            params.load_mode = LLAMA_LOAD_MODE_MMAP
             return params
         }
 
@@ -829,6 +828,7 @@ import Foundation
                 llama_sampler_chain_add(
                     samplerPtr,
                     llama_sampler_init_penalties(
+                        llama_vocab_n_tokens(vocab),
                         effectiveRepeatLastN,
                         effectiveRepeatPenalty,
                         effectiveFrequencyPenalty,
@@ -958,6 +958,7 @@ import Foundation
                 llama_sampler_chain_add(
                     samplerPointer,
                     llama_sampler_init_penalties(
+                        llama_vocab_n_tokens(vocab),
                         options.repeatLastN,
                         options.repeatPenalty,
                         options.frequencyPenalty,
@@ -1109,18 +1110,37 @@ import Foundation
             }
 
             mutating func sample(from allowedTokens: Set<Int>) async throws -> Int {
-                guard let logits = llama_get_logits(context) else {
+                // Masking llama_get_logits in place does not constrain
+                // llama_sampler_sample: the context can refresh that buffer when
+                // the sampler fetches logits, dropping the mask. Build a candidate
+                // array holding only the allowed tokens and run the sampler chain
+                // over it instead.
+                guard let logits = llama_get_logits_ith(context, batch.pointee.n_tokens - 1) else {
                     return eosToken
                 }
 
-                for tokenIndex in 0 ..< vocabSize {
-                    if !allowedTokens.contains(tokenIndex) {
-                        logits[tokenIndex] = -Float.infinity
-                    }
+                var candidates = allowedTokens.compactMap { token -> llama_token_data? in
+                    guard token >= 0, token < vocabSize else { return nil }
+                    return llama_token_data(id: llama_token(token), logit: logits[token], p: 0)
+                }
+                guard !candidates.isEmpty else {
+                    return eosToken
                 }
 
-                let tokenIndex = batch.pointee.n_tokens - 1
-                return Int(llama_sampler_sample(sampler, context, tokenIndex))
+                let selectedToken = candidates.withUnsafeMutableBufferPointer { buffer -> Int in
+                    var array = llama_token_data_array(
+                        data: buffer.baseAddress,
+                        size: buffer.count,
+                        selected: -1,
+                        sorted: false
+                    )
+                    llama_sampler_apply(sampler, &array)
+                    guard array.selected >= 0, array.selected < Int64(array.size), let data = array.data else {
+                        return eosToken
+                    }
+                    return Int(data[Int(array.selected)].id)
+                }
+                return selectedToken
             }
         }
 
@@ -1197,6 +1217,7 @@ import Foundation
                     llama_sampler_chain_add(
                         samplerPtr,
                         llama_sampler_init_penalties(
+                            llama_vocab_n_tokens(vocab),
                             effectiveRepeatLastN,
                             effectiveRepeatPenalty,
                             effectiveFrequencyPenalty,