diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index d2f29b03..e6204a1f 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -208,6 +208,11 @@ import Foundation /// A negative value offloads all layers, and `0` runs entirely on the CPU. public let gpuLayers: Int32 + /// The path to the multimodal projector GGUF file, when the model has one. + /// + /// Prompts may include image segments only when a projector is loaded. + public let mmprojPath: String? + /// The default GPU layer count for the current platform. /// /// All layers are offloaded by default: the prebuilt llama.cpp binaries @@ -453,18 +458,258 @@ import Foundation /// The model's vocabulary private var vocab: OpaquePointer? + /// The multimodal projector context, when a projector file was provided + private var mtmdContext: OpaquePointer? + /// Whether the model is currently loaded private var isModelLoaded: Bool = false + /// A context kept alive for one session so exchanges reuse its state. + private struct CachedSessionContext { + let sessionID: ObjectIdentifier + let context: OpaquePointer + var tokens: [llama_token] + let contextSize: UInt32 + let batchSize: UInt32 + /// Whether a generation is currently decoding on the context. + var isCheckedOut: Bool + /// Whether the context should be freed once the current generation releases it. + var discardWhenReleased: Bool + } + + /// Guards `cachedSessionContext`. A generation checks the cached context out + /// for its whole run, so a concurrent generation for another session never + /// frees a context that is still decoding: it runs on a transient context + /// instead and leaves the cache untouched. + private let sessionContextLock = NSLock() + private var cachedSessionContext: CachedSessionContext? + + /// The number of prompt tokens reused from the cached context by the most + /// recent chat generation. + internal private(set) var lastReusedTokenCount: Int = 0 + + /// The number of prompt tokens decoded by the most recent chat generation. + internal private(set) var lastPrefillTokenCount: Int = 0 + + /// Frees the cached per-session context and the state it holds. + /// + /// The cached context, including its KV state, otherwise lives as long as the + /// model. The next chat generation prefills its full prompt again. Call this + /// under memory pressure or when a session is discarded. If a generation is + /// running on the cached context, it is freed as soon as that generation ends. + public func clearCachedContext() { + discardCachedSessionContext() + } + + private func discardCachedSessionContext() { + sessionContextLock.lock() + defer { sessionContextLock.unlock() } + guard var cached = cachedSessionContext else { return } + if cached.isCheckedOut { + cached.discardWhenReleased = true + cachedSessionContext = cached + return + } + llama_free(cached.context) + cachedSessionContext = nil + } + + private func recordCachedTokens(_ tokens: [llama_token], context: OpaquePointer) { + sessionContextLock.lock() + defer { sessionContextLock.unlock() } + guard var cached = cachedSessionContext, cached.context == context else { return } + cached.tokens = tokens + cachedSessionContext = cached + } + + /// Returns a context obtained from `acquireSessionContext`. The cached + /// context is checked back in (or freed, if a discard was requested while + /// it was busy); a transient context is freed. + private func releaseSessionContext(_ context: OpaquePointer) { + sessionContextLock.lock() + defer { sessionContextLock.unlock() } + if var cached = cachedSessionContext, cached.context == context { + if cached.discardWhenReleased { + llama_free(cached.context) + cachedSessionContext = nil + } else { + cached.isCheckedOut = false + cachedSessionContext = cached + } + return + } + llama_free(context) + } + + /// Returns a context for the session along with the index of the first + /// prompt token that still needs to be decoded. Pair every call with + /// `releaseSessionContext(_:)` once generation ends. + /// + /// A cached context whose recorded tokens share a prefix with the prompt + /// keeps that prefix: matching state past the divergence point is removed + /// with `llama_memory_seq_rm`, and backends whose state cannot be rewound + /// (recurrent models) fall back to clearing the memory and decoding the + /// full prompt. The final prompt token is always re-decoded so sampling + /// has fresh logits. + /// + /// While another generation holds the cached context, the caller gets a + /// transient context that decodes the full prompt and is not cached. + private func acquireSessionContext( + for session: LanguageModelSession, + promptTokens: [llama_token], + options: ResolvedGenerationOptions + ) throws -> (context: OpaquePointer, startIndex: Int) { + let sessionID = ObjectIdentifier(session) + sessionContextLock.lock() + defer { sessionContextLock.unlock() } + + if var cached = cachedSessionContext, + !cached.isCheckedOut, + cached.sessionID == sessionID, + cached.contextSize == options.contextSize, + cached.batchSize == options.batchSize + { + var common = 0 + while common < cached.tokens.count, common < promptTokens.count, + cached.tokens[common] == promptTokens[common] + { + common += 1 + } + if common == promptTokens.count { + common = max(0, promptTokens.count - 1) + } + if common < cached.tokens.count { + let memory = llama_get_memory(cached.context) + if !llama_memory_seq_rm(memory, 0, llama_pos(common), -1) { + llama_memory_clear(memory, true) + common = 0 + } + } + cached.tokens = Array(promptTokens.prefix(common)) + cached.isCheckedOut = true + cachedSessionContext = cached + return (cached.context, common) + } + + if let cached = cachedSessionContext, cached.isCheckedOut { + return (try makeFreshContext(options: options), 0) + } + + if let cached = cachedSessionContext { + llama_free(cached.context) + cachedSessionContext = nil + } + let contextParams = createContextParams(from: options) + guard let context = llama_init_from_model(model!, contextParams) else { + throw LlamaLanguageModelError.contextInitializationFailed + } + guard llama_get_memory(context) != nil else { + llama_free(context) + throw LlamaLanguageModelError.encoderOnlyModel + } + cachedSessionContext = CachedSessionContext( + sessionID: sessionID, + context: context, + tokens: [], + contextSize: options.contextSize, + batchSize: options.batchSize, + isCheckedOut: true, + discardWhenReleased: false + ) + return (context, 0) + } + + /// Creates a single-use context for generations that do not reuse state. + private func makeFreshContext(options: ResolvedGenerationOptions) throws -> OpaquePointer { + let contextParams = createContextParams(from: options) + guard let context = llama_init_from_model(model!, contextParams) else { + throw LlamaLanguageModelError.contextInitializationFailed + } + guard llama_get_memory(context) != nil else { + llama_free(context) + throw LlamaLanguageModelError.encoderOnlyModel + } + llama_set_causal_attn(context, true) + llama_set_n_threads(context, options.threads, options.threads) + return context + } + + /// Runs a chat text generation for the session, reusing the session's + /// cached context when its state matches a prefix of the prompt. + private func generateChatText( + session: LanguageModelSession, + prompt: String, + maxTokens: Int, + options: ResolvedGenerationOptions, + onToken: (String) -> Void + ) throws { + guard let model = self.model, let vocab = llama_model_get_vocab(model) else { + throw LlamaLanguageModelError.contextInitializationFailed + } + + let promptTokens = try tokenizeText(vocab: vocab, text: prompt) + guard !promptTokens.isEmpty else { + throw LlamaLanguageModelError.tokenizationFailed + } + + if llama_model_has_encoder(model) { + let context = try makeFreshContext(options: options) + defer { llama_free(context) } + try performTokenGeneration( + context: context, + vocab: vocab, + promptTokens: promptTokens, + startIndex: 0, + maxTokens: maxTokens, + options: options, + onToken: onToken + ) + return + } + + let (context, startIndex) = try acquireSessionContext( + for: session, + promptTokens: promptTokens, + options: options + ) + defer { releaseSessionContext(context) } + llama_set_causal_attn(context, true) + llama_set_n_threads(context, options.threads, options.threads) + + do { + try performTokenGeneration( + context: context, + vocab: vocab, + promptTokens: promptTokens, + startIndex: startIndex, + maxTokens: maxTokens, + options: options, + onToken: onToken + ) + } catch { + discardCachedSessionContext() + throw error + } + } + /// Creates a Llama language model. /// /// - Parameters: /// - modelPath: The path to the GGUF model file. /// - gpuLayers: The number of model layers to offload to the GPU. /// Defaults to ``defaultGPULayerCount``. - public init(modelPath: String, gpuLayers: Int32 = LlamaLanguageModel.defaultGPULayerCount) { + /// - mmprojPath: The path to a multimodal projector GGUF file matching + /// the model. When provided, prompts may include image segments, + /// which are encoded through the projector. Defaults to `nil` + /// (text only). + public init( + modelPath: String, + gpuLayers: Int32 = LlamaLanguageModel.defaultGPULayerCount, + mmprojPath: String? = nil + ) { self.modelPath = modelPath self.gpuLayers = gpuLayers + self.mmprojPath = mmprojPath self.legacyDefaults = ResolvedGenerationOptions() } @@ -505,6 +750,12 @@ import Foundation } deinit { + if let cached = cachedSessionContext { + llama_free(cached.context) + } + if let mtmdContext = mtmdContext { + mtmd_free(mtmdContext) + } if let model = model { llama_model_free(model) } @@ -517,30 +768,16 @@ import Foundation includeSchemaInPrompt: Bool, options: GenerationOptions ) async throws -> LanguageModelSession.Response where Content: Generable { - // Validate that no image segments are present - try validateNoImageSegments(in: session) + if mmprojPath == nil { + try validateNoImageSegments(in: session) + } try await ensureModelLoaded() let runtimeOptions = resolvedOptions(from: options) let structuredOptions = resolvedStructuredOptions(from: options) - let contextParams = createContextParams(from: runtimeOptions) - - // Try to create context with error handling - guard let context = llama_init_from_model(model!, contextParams) else { - throw LlamaLanguageModelError.contextInitializationFailed - } - defer { llama_free(context) } - // Check if this is an embedding model (no KV cache). - // This early check catches models configured for embeddings that lack a KV cache. - // A complementary architectural check in prepareInitialBatch catches encoder-only - // models (like BERT) by their architecture type. - if llama_get_memory(context) == nil { - throw LlamaLanguageModelError.encoderOnlyModel - } - - llama_set_causal_attn(context, true) - llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) + var promptImages: [Data] = [] + let imageMarker = mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil let fullPrompt: String if includeSchemaInPrompt, type != String.self { @@ -550,18 +787,45 @@ import Foundation assistantPrefill: runtimeOptions.assistantPrefill ) } else { - fullPrompt = try formatPrompt(for: session, assistantPrefill: runtimeOptions.assistantPrefill) + fullPrompt = try formatPrompt( + for: session, + extraSystemMessage: nil, + assistantPrefill: runtimeOptions.assistantPrefill, + imageMarker: imageMarker, + images: &promptImages + ) } if type == String.self { let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 - let text = try await generateText( - context: context, - model: model!, - prompt: fullPrompt, - maxTokens: maxTokens, - options: runtimeOptions - ) + let text: String + if promptImages.isEmpty { + var accumulated = "" + try generateChatText( + session: session, + prompt: fullPrompt, + maxTokens: maxTokens, + options: runtimeOptions + ) { tokenText in + accumulated += tokenText + } + text = accumulated + } else { + discardCachedSessionContext() + let context = try makeFreshContext(options: runtimeOptions) + defer { llama_free(context) } + var accumulated = "" + try performMultimodalGeneration( + context: context, + prompt: fullPrompt, + images: promptImages, + maxTokens: maxTokens, + options: runtimeOptions + ) { tokenText in + accumulated += tokenText + } + text = accumulated + } return LanguageModelSession.Response( content: text as! Content, @@ -569,6 +833,8 @@ import Foundation transcriptEntries: ArraySlice([]) ) } else { + let context = try makeFreshContext(options: runtimeOptions) + defer { llama_free(context) } let maxTokens = structuredOptions.maximumResponseTokens ?? 512 let jsonString = try await generateStructuredJSON( context: context, @@ -599,15 +865,16 @@ import Foundation fatalError("LlamaLanguageModel only supports generating String content") } - // Validate that no image segments are present - do { - try validateNoImageSegments(in: session) - } catch { - return LanguageModelSession.ResponseStream( - stream: AsyncThrowingStream { continuation in - continuation.finish(throwing: error) - } - ) + if mmprojPath == nil { + do { + try validateNoImageSegments(in: session) + } catch { + return LanguageModelSession.ResponseStream( + stream: AsyncThrowingStream { continuation in + continuation.finish(throwing: error) + } + ) + } } let stream: AsyncThrowingStream.Snapshot, any Error> = @@ -618,37 +885,20 @@ import Foundation let runtimeOptions = resolvedOptions(from: options) let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 - let contextParams = createContextParams(from: runtimeOptions) - guard let context = llama_init_from_model(model!, contextParams) else { - throw LlamaLanguageModelError.contextInitializationFailed - } - defer { llama_free(context) } - - // Check if this is an embedding model (no KV cache). - // This early check catches models configured for embeddings that lack a KV cache. - // A complementary architectural check in prepareInitialBatch catches encoder-only - // models (like BERT) by their architecture type. - if llama_get_memory(context) == nil { - throw LlamaLanguageModelError.encoderOnlyModel - } - - // Stabilize runtime behavior per-context - llama_set_causal_attn(context, true) - llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) var accumulatedText = "" + var promptImages: [Data] = [] + let imageMarker = + self.mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil let fullPrompt = try self.formatPrompt( for: session, - assistantPrefill: runtimeOptions.assistantPrefill + extraSystemMessage: nil, + assistantPrefill: runtimeOptions.assistantPrefill, + imageMarker: imageMarker, + images: &promptImages ) - try self.performTextGeneration( - context: context, - model: model!, - prompt: fullPrompt, - maxTokens: maxTokens, - options: runtimeOptions - ) { tokenText in + let yieldToken: (String) -> Void = { tokenText in accumulatedText += tokenText let snapshot = LanguageModelSession.ResponseStream.Snapshot( @@ -658,6 +908,28 @@ import Foundation continuation.yield(snapshot) } + if promptImages.isEmpty { + try self.generateChatText( + session: session, + prompt: fullPrompt, + maxTokens: maxTokens, + options: runtimeOptions, + onToken: yieldToken + ) + } else { + self.discardCachedSessionContext() + let context = try self.makeFreshContext(options: runtimeOptions) + defer { llama_free(context) } + try self.performMultimodalGeneration( + context: context, + prompt: fullPrompt, + images: promptImages, + maxTokens: maxTokens, + options: runtimeOptions, + onToken: yieldToken + ) + } + continuation.finish() } catch { continuation.finish(throwing: error) @@ -686,6 +958,11 @@ import Foundation llama_backend_init() // Free any existing model before loading a new one + discardCachedSessionContext() + if let existingContext = mtmdContext { + mtmd_free(existingContext) + self.mtmdContext = nil + } if let existingModel = model { llama_model_free(existingModel) self.model = nil @@ -696,6 +973,22 @@ import Foundation throw LlamaLanguageModelError.modelLoadFailed } + if let mmprojPath { + guard FileManager.default.fileExists(atPath: mmprojPath) else { + llama_model_free(loadedModel) + throw LlamaLanguageModelError.invalidModelPath + } + var mtmdParams = mtmd_context_params_default() + mtmdParams.use_gpu = gpuLayers != 0 + mtmdParams.print_timings = false + mtmdParams.n_threads = legacyDefaults.threads + guard let projector = mtmd_init_from_file(mmprojPath, loadedModel, mtmdParams) else { + llama_model_free(loadedModel) + throw LlamaLanguageModelError.modelLoadFailed + } + self.mtmdContext = projector + } + self.model = loadedModel self.vocab = llama_model_get_vocab(loadedModel) self.isModelLoaded = true @@ -811,110 +1104,6 @@ import Foundation llama_sampler_chain_add(sampler, llama_sampler_init_dist(options.seed)) } - private func generateText( - context: OpaquePointer, - model: OpaquePointer, - prompt: String, - maxTokens: Int, - options: ResolvedGenerationOptions - ) async throws - -> String - { - guard let vocab = llama_model_get_vocab(model) else { - throw LlamaLanguageModelError.contextInitializationFailed - } - - // Tokenize the prompt - let promptTokens = try tokenizeText(vocab: vocab, text: prompt) - guard !promptTokens.isEmpty else { - throw LlamaLanguageModelError.tokenizationFailed - } - - var batch = llama_batch_init(Int32(options.batchSize), 0, 1) - defer { llama_batch_free(batch) } - - let hasEncoder = try prepareInitialBatch( - batch: &batch, - promptTokens: promptTokens, - model: model, - vocab: vocab, - context: context, - batchSize: options.batchSize, - contextSize: options.contextSize - ) - - // Initialize sampler chain with options - guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { - throw LlamaLanguageModelError.decodingFailed - } - defer { llama_sampler_free(sampler) } - let samplerPtr = UnsafeMutablePointer(sampler) - - let effectiveTemperature = Float(options.temperature) - - // Apply repeat/frequency/presence penalties from custom options - let effectiveRepeatPenalty = options.repeatPenalty - let effectiveRepeatLastN = options.repeatLastN - let effectiveFrequencyPenalty = options.frequencyPenalty - let effectivePresencePenalty = options.presencePenalty - - if effectiveRepeatPenalty != 1.0 || effectiveFrequencyPenalty != 0.0 || effectivePresencePenalty != 0.0 { - llama_sampler_chain_add( - samplerPtr, - llama_sampler_init_penalties( - llama_vocab_n_tokens(vocab), - effectiveRepeatLastN, - effectiveRepeatPenalty, - effectiveFrequencyPenalty, - effectivePresencePenalty - ) - ) - } - - applySampling(sampler: samplerPtr, effectiveTemperature: effectiveTemperature, options: options) - - // Generate tokens one by one - var generatedText = "" - // Track position - for encoder-decoder models, we start from position 1 (after decoder start token) - // For decoder-only models, we continue from the end of the prompt - var n_cur: Int32 = hasEncoder ? 1 : Int32(promptTokens.count) - - for _ in 0 ..< maxTokens { - // Sample next token from logits - llama_batch_get_one creates batch with single token at index 0 - let nextToken = llama_sampler_sample(sampler, context, batch.n_tokens - 1) - llama_sampler_accept(sampler, nextToken) - - // Check for end of sequence - if llama_vocab_is_eog(vocab, nextToken) { - break - } - - // Convert token to text - if let tokenText = tokenToText(vocab: vocab, token: nextToken) { - generatedText += tokenText - } - - // Prepare batch for next token - batch.n_tokens = 1 - batch.token[0] = nextToken - batch.pos[0] = n_cur - batch.n_seq_id[0] = 1 - if let seq_ids = batch.seq_id, let seq_id = seq_ids[0] { - seq_id[0] = 0 - } - batch.logits[0] = 1 - - n_cur += 1 - - let decodeResult = llama_decode(context, batch) - guard decodeResult == 0 else { - break - } - } - - return generatedText - } - /// Builds a JSONSchema-informed prompt for structured output. private func schemaPrompt(for schema: GenerationSchema) -> String { let encoder = JSONEncoder() @@ -1179,23 +1368,21 @@ import Foundation } } - private func performTextGeneration( + private func performTokenGeneration( context: OpaquePointer, - model: OpaquePointer, - prompt: String, + vocab: OpaquePointer, + promptTokens: [llama_token], + startIndex: Int, maxTokens: Int, options: ResolvedGenerationOptions, onToken: (String) -> Void ) throws { - guard let vocab = llama_model_get_vocab(model) else { - throw LlamaLanguageModelError.contextInitializationFailed + guard let model = self.model else { + throw LlamaLanguageModelError.modelLoadFailed } - // Tokenize the prompt - let promptTokens = try tokenizeText(vocab: vocab, text: prompt) - guard !promptTokens.isEmpty else { - throw LlamaLanguageModelError.tokenizationFailed - } + lastReusedTokenCount = startIndex + lastPrefillTokenCount = promptTokens.count - startIndex // Initialize batch var batch = llama_batch_init(Int32(options.batchSize), 0, 1) @@ -1208,7 +1395,8 @@ import Foundation vocab: vocab, context: context, batchSize: options.batchSize, - contextSize: options.contextSize + contextSize: options.contextSize, + startIndex: startIndex ) // Initialize sampler chain with options @@ -1246,6 +1434,7 @@ import Foundation // Track position - for encoder-decoder models, we start from position 1 (after decoder start token) // For decoder-only models, we continue from the end of the prompt var n_cur: Int32 = hasEncoder ? 1 : Int32(promptTokens.count) + var decodedTokens = promptTokens for _ in 0 ..< maxTokens { if Task.isCancelled { @@ -1282,6 +1471,149 @@ import Foundation guard decodeResult == 0 else { break } + decodedTokens.append(nextToken) + } + + recordCachedTokens(decodedTokens, context: context) + } + + /// Evaluates a marker-annotated multimodal prompt through the projector, + /// then generates text tokens from the resulting state. + private func performMultimodalGeneration( + context: OpaquePointer, + prompt: String, + images: [Data], + maxTokens: Int, + options: ResolvedGenerationOptions, + onToken: (String) -> Void + ) throws { + guard let mtmdContext, let model = self.model, + let vocab = llama_model_get_vocab(model) + else { + throw LlamaLanguageModelError.contextInitializationFailed + } + + var bitmaps: [OpaquePointer?] = [] + defer { + for bitmap in bitmaps { + if let bitmap { + mtmd_bitmap_free(bitmap) + } + } + } + for imageData in images { + // Pinned to the current llama.swift signature. llama.cpp master adds a + // trailing options argument to this helper; update alongside the dependency. + let wrapper = imageData.withUnsafeBytes { raw -> mtmd_helper_bitmap_wrapper in + mtmd_helper_bitmap_init_from_buf( + mtmdContext, + raw.bindMemory(to: UInt8.self).baseAddress, + imageData.count, + false + ) + } + if let videoContext = wrapper.video_ctx { + mtmd_helper_video_free(videoContext) + throw LlamaLanguageModelError.unsupportedFeature + } + guard let bitmap = wrapper.bitmap else { + throw LlamaLanguageModelError.encodingFailed + } + bitmaps.append(bitmap) + } + + guard let chunks = mtmd_input_chunks_init() else { + throw LlamaLanguageModelError.encodingFailed + } + defer { mtmd_input_chunks_free(chunks) } + + let tokenizeResult = prompt.withCString { cPrompt -> Int32 in + var inputText = mtmd_input_text( + text: cPrompt, + text_len: strlen(cPrompt), + add_special: true, + parse_special: true + ) + return bitmaps.withUnsafeMutableBufferPointer { buffer in + mtmd_tokenize(mtmdContext, chunks, &inputText, buffer.baseAddress, buffer.count) + } + } + guard tokenizeResult == 0 else { + throw LlamaLanguageModelError.tokenizationFailed + } + + var pastPosition: llama_pos = 0 + let evalResult = mtmd_helper_eval_chunks( + mtmdContext, + context, + chunks, + 0, + 0, + Int32(options.batchSize), + true, + &pastPosition + ) + guard evalResult == 0 else { + throw LlamaLanguageModelError.decodingFailed + } + + guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { + throw LlamaLanguageModelError.decodingFailed + } + defer { llama_sampler_free(sampler) } + let samplerPtr = UnsafeMutablePointer(sampler) + + if options.repeatPenalty != 1.0 || options.frequencyPenalty != 0.0 || options.presencePenalty != 0.0 { + llama_sampler_chain_add( + samplerPtr, + llama_sampler_init_penalties( + llama_vocab_n_tokens(vocab), + options.repeatLastN, + options.repeatPenalty, + options.frequencyPenalty, + options.presencePenalty + ) + ) + } + applySampling(sampler: samplerPtr, effectiveTemperature: options.temperature, options: options) + + var batch = llama_batch_init(1, 0, 1) + defer { llama_batch_free(batch) } + + var n_cur: Int32 = Int32(pastPosition) + var sampleIndex: Int32 = -1 + + for _ in 0 ..< maxTokens { + if Task.isCancelled { + break + } + + let nextToken = llama_sampler_sample(samplerPtr, context, sampleIndex) + llama_sampler_accept(samplerPtr, nextToken) + + if llama_vocab_is_eog(vocab, nextToken) { + break + } + + if let tokenText = tokenToText(vocab: vocab, token: nextToken) { + onToken(tokenText) + } + + batch.n_tokens = 1 + batch.token[0] = nextToken + batch.pos[0] = n_cur + batch.n_seq_id[0] = 1 + if let seq_ids = batch.seq_id, let seq_id = seq_ids[0] { + seq_id[0] = 0 + } + batch.logits[0] = 1 + + n_cur += 1 + + guard llama_decode(context, batch) == 0 else { + break + } + sampleIndex = 0 } } @@ -1316,6 +1648,8 @@ import Foundation /// - context: The model context. /// - batchSize: The batch capacity per decode call. /// - contextSize: The context window the prompt must fit within. + /// - startIndex: The index of the first prompt token to decode. Earlier + /// tokens are already present in the context's state. Defaults to `0`. /// - Returns: `true` if the model has an encoder (for position tracking during generation). /// - Throws: `promptExceedsContextWindow` if the prompt cannot fit in the context window, /// `insufficientMemory` if an encoder prompt exceeds the batch capacity, `encoderOnlyModel` @@ -1327,7 +1661,8 @@ import Foundation vocab: OpaquePointer, context: OpaquePointer, batchSize: UInt32, - contextSize: UInt32 + contextSize: UInt32, + startIndex: Int = 0 ) throws -> Bool { // Leave at least one context cell free for generation. guard promptTokens.count < contextSize else { @@ -1338,7 +1673,7 @@ import Foundation let hasDecoder = llama_model_has_decoder(model) // Encoder models ingest the full prompt in a single llama_encode call. - guard !hasEncoder || promptTokens.count <= batchSize else { + guard !hasEncoder || (startIndex == 0 && promptTokens.count <= batchSize) else { throw LlamaLanguageModelError.insufficientMemory } @@ -1390,7 +1725,7 @@ import Foundation // batch-sized chunks with absolute positions, requesting logits // only for the final token. let capacity = Int(batchSize) - var start = 0 + var start = startIndex while start < promptTokens.count { let count = min(capacity, promptTokens.count - start) batch.n_tokens = Int32(count) @@ -1422,6 +1757,23 @@ import Foundation for session: LanguageModelSession, extraSystemMessage: String? = nil, assistantPrefill: String? = nil + ) throws -> String { + var images: [Data] = [] + return try formatPrompt( + for: session, + extraSystemMessage: extraSystemMessage, + assistantPrefill: assistantPrefill, + imageMarker: nil, + images: &images + ) + } + + private func formatPrompt( + for session: LanguageModelSession, + extraSystemMessage: String?, + assistantPrefill: String?, + imageMarker: String?, + images: inout [Data] ) throws -> String { guard let model = self.model else { throw LlamaLanguageModelError.modelLoadFailed @@ -1432,19 +1784,31 @@ import Foundation for entry in session.transcript { switch entry { case .instructions(let instructions): - let text = extractText(from: instructions.segments) + let text = try extractContent( + from: instructions.segments, + imageMarker: imageMarker, + images: &images + ) if !text.isEmpty { messages.append(("system", text)) } case .prompt(let prompt): - let text = extractText(from: prompt.segments) + let text = try extractContent( + from: prompt.segments, + imageMarker: imageMarker, + images: &images + ) if !text.isEmpty { messages.append(("user", text)) } case .response(let response): - let text = extractText(from: response.segments) + let text = try extractContent( + from: response.segments, + imageMarker: imageMarker, + images: &images + ) if !text.isEmpty { messages.append(("assistant", text)) } @@ -1486,6 +1850,9 @@ import Foundation ) guard requiredSize > 0 else { + if let tmpl, String(cString: tmpl).contains("<|turn>") { + return Self.renderGemma4Prompt(messages: messages, assistantPrefill: assistantPrefill) + } throw LlamaLanguageModelError.encodingFailed } @@ -1515,6 +1882,27 @@ import Foundation return rendered } + /// Renders the Gemma 4 canonical chat format, which + /// `llama_chat_apply_template` does not recognize: turns open with + /// `<|turn>role`, close with ``, and the assistant role is named + /// `model`. The BOS token is applied during tokenization. + static func renderGemma4Prompt( + messages: [(role: String, content: String)], + assistantPrefill: String? + ) -> String { + var rendered = "" + for message in messages { + let role = message.role == "assistant" ? "model" : message.role + let content = message.content.trimmingCharacters(in: .whitespacesAndNewlines) + rendered += "<|turn>\(role)\n\(content)\n" + } + rendered += "<|turn>model\n" + if let assistantPrefill, !assistantPrefill.isEmpty { + rendered += assistantPrefill + } + return rendered + } + private func extractText(from segments: [Transcript.Segment]) -> String { segments.compactMap { segment -> String? in if case .text(let t) = segment { return t.content } @@ -1522,6 +1910,42 @@ import Foundation }.joined() } + /// Extracts message content from segments, replacing each image segment + /// with `imageMarker` and collecting its payload in order. Image segments + /// throw ``LlamaLanguageModelError/unsupportedFeature`` when no marker is + /// provided. + private func extractContent( + from segments: [Transcript.Segment], + imageMarker: String?, + images: inout [Data] + ) throws -> String { + var parts: [String] = [] + for segment in segments { + switch segment { + case .text(let t): + parts.append(t.content) + case .image(let image): + guard let imageMarker else { + throw LlamaLanguageModelError.unsupportedFeature + } + switch image.source { + case .data(let data, _): + images.append(data) + parts.append(imageMarker) + case .url(let url): + guard url.isFileURL, let data = try? Data(contentsOf: url) else { + throw LlamaLanguageModelError.unsupportedFeature + } + images.append(data) + parts.append(imageMarker) + } + default: + break + } + } + return parts.joined() + } + private func tokenizeText(vocab: OpaquePointer, text: String) throws -> [llama_token] { let utf8Count = text.utf8.count let maxTokens = Int32(max(utf8Count * 2, 8)) // Rough estimate, minimum capacity diff --git a/Tests/AnyLanguageModelTests/LlamaGemma4TemplateTests.swift b/Tests/AnyLanguageModelTests/LlamaGemma4TemplateTests.swift new file mode 100644 index 00000000..756eb412 --- /dev/null +++ b/Tests/AnyLanguageModelTests/LlamaGemma4TemplateTests.swift @@ -0,0 +1,37 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +#if Llama + @Suite("LlamaLanguageModel Gemma 4 template") + struct LlamaGemma4TemplateTests { + @Test func rendersTurnsAndOpensTheModelTurn() { + let rendered = LlamaLanguageModel.renderGemma4Prompt( + messages: [ + (role: "system", content: "Be brief."), + (role: "user", content: "Hi"), + (role: "assistant", content: "Hello"), + (role: "user", content: "Bye"), + ], + assistantPrefill: nil + ) + #expect( + rendered + == "<|turn>system\nBe brief.\n" + + "<|turn>user\nHi\n" + + "<|turn>model\nHello\n" + + "<|turn>user\nBye\n" + + "<|turn>model\n" + ) + } + + @Test func appendsTheAssistantPrefill() { + let rendered = LlamaLanguageModel.renderGemma4Prompt( + messages: [(role: "user", content: "Hi")], + assistantPrefill: "" + ) + #expect(rendered == "<|turn>user\nHi\n<|turn>model\n") + } + } +#endif diff --git a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift index a319b82c..e1aa4640 100644 --- a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift @@ -47,6 +47,26 @@ import Testing #expect(!response.content.isEmpty) } + @Test func reusesSessionContextAcrossTurns() async throws { + let session = LanguageModelSession(model: model) + var options = GenerationOptions(maximumResponseTokens: 24) + options[custom: LlamaLanguageModel.self] = .init(contextSize: 2048, batchSize: 512) + + let first = try await session.respond( + to: "My favorite color is blue. Reply with OK.", + options: options + ) + #expect(!first.content.isEmpty) + #expect(model.lastReusedTokenCount == 0) + + let second = try await session.respond( + to: "What is my favorite color? Answer with one word.", + options: options + ) + #expect(!second.content.isEmpty) + #expect(model.lastReusedTokenCount > 0) + } + @Test func customGenerationOptionsRoundTrip() { var options = GenerationOptions( temperature: 0.6, @@ -425,3 +445,78 @@ import Testing } } #endif // Llama + +#if Llama + @Suite( + "LlamaLanguageModel vision", + .serialized, + .enabled( + if: ProcessInfo.processInfo.environment["LLAMA_VISION_MODEL_PATH"] != nil + && ProcessInfo.processInfo.environment["LLAMA_VISION_MMPROJ_PATH"] != nil + ) + ) + struct LlamaLanguageModelVisionTests { + static let redSquarePNG = Data( + base64Encoded: "iVBORw0KGgoAAAANSUhEUgAAAGAAAABgCAIAAABt+uBvAAABC0lEQVR4nO3OMQ0AIAAEsfdvGhyw9gaS" + + "CujO9j34QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UHcBWwZ3g5gacwjAAAAAElFTkSuQmCC" + )! + + let model = LlamaLanguageModel( + modelPath: ProcessInfo.processInfo.environment["LLAMA_VISION_MODEL_PATH"]!, + mmprojPath: ProcessInfo.processInfo.environment["LLAMA_VISION_MMPROJ_PATH"]! + ) + + @Test func describesImageData() async throws { + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt(segments: [ + .text(.init(content: "What is the dominant color of this image? Answer with one word.")), + .image(.init(data: Self.redSquarePNG, mimeType: "image/png")), + ]) + ) + ]) + let session = LanguageModelSession(model: model, transcript: transcript) + let response = try await session.respond(to: "") + #expect(response.content.lowercased().contains("red")) + } + + @Test func streamsImageDescription() async throws { + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt(segments: [ + .text(.init(content: "What is the dominant color of this image? Answer with one word.")), + .image(.init(data: Self.redSquarePNG, mimeType: "image/png")), + ]) + ) + ]) + let session = LanguageModelSession(model: model, transcript: transcript) + let stream = session.streamResponse(to: "") + var last = "" + for try await snapshot in stream { + last = snapshot.content + } + #expect(last.lowercased().contains("red")) + } + + @Test func rejectsImagesWithoutProjector() async throws { + let textOnlyModel = LlamaLanguageModel( + modelPath: ProcessInfo.processInfo.environment["LLAMA_VISION_MODEL_PATH"]! + ) + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt(segments: [ + .image(.init(data: Self.redSquarePNG, mimeType: "image/png")) + ]) + ) + ]) + let session = LanguageModelSession(model: textOnlyModel, transcript: transcript) + await #expect(throws: LlamaLanguageModelError.self) { + _ = try await session.respond(to: "") + } + } + } +#endif