markbaseengine/Sources/MarkBase/Generator/StreamingGenerator.swift

import Foundation

// ─────────────────────────────────────────────────────────────
// Generation Configuration
// ─────────────────────────────────────────────────────────────

public struct GenerationConfig: Sendable {
    public let maxTokens: Int
    public let temperature: Float
    public let topK: Int?
    public let topP: Float?
    public let stopTokens: [Int]?
    
    public init(
        maxTokens: Int = 100,
        temperature: Float = 1.0,
        topK: Int? = nil,
        topP: Float? = nil,
        stopTokens: [Int]? = nil
    ) {
        self.maxTokens = maxTokens
        self.temperature = temperature
        self.topK = topK
        self.topP = topP
        self.stopTokens = stopTokens
    }
    
    // Default configuration
    public static let defaultConfig = GenerationConfig(maxTokens: 100, temperature: 1.0)
}

// ─────────────────────────────────────────────────────────────
// Streaming Generator - Token-by-token generation
// ─────────────────────────────────────────────────────────────

public final class StreamingGenerator: @unchecked Sendable {
    private let model: E4BModel
    private let tokenizer: Tokenizer
    private let engine: MarkBaseEngine
    private let sampler: Sampler
    
    public init(model: E4BModel, tokenizer: Tokenizer, engine: MarkBaseEngine) {
        self.model = model
        self.tokenizer = tokenizer
        self.engine = engine
        self.sampler = Sampler()
    }
    
    // ─────────────────────────────────────────────────────────────
    // Stream Generate - AsyncStream for token-by-token output
    // ─────────────────────────────────────────────────────────────
    
    public func generate(
        prompt: String,
        config: GenerationConfig = .defaultConfig
    ) -> AsyncStream<String> {
        return AsyncStream { continuation in
            Task {
                do {
                    // Encode prompt
                    let promptTokens = tokenizer.encode(text: prompt)
                    
                    // Pre-fill KV cache with prompt tokens
                    var lastLogits: [Float] = []
                    for (position, tokenId) in promptTokens.enumerated() {
                        lastLogits = try model.forward(tokenId: tokenId, position: position)
                    }
                    
                    // Generate tokens
                    var generatedTokens: [Int] = []
                    var position = promptTokens.count
                    var streamDecoder = StreamingDecoder(tokenizer: tokenizer)
                    
                    for _ in 0..<config.maxTokens {
                        // Sample next token
                        let nextToken = sampler.sample(
                            logits: lastLogits,
                            temperature: config.temperature,
                            topK: config.topK,
                            topP: config.topP
                        )
                        
                        // Debug: print selected token
                        if generatedTokens.count < 5 {
                            print("[DEBUG] Generated token \(generatedTokens.count): ID=\(nextToken), raw='\(tokenizer.rawToken(for: nextToken) ?? "nil")'")
                            fflush(stdout)
                        }
                        
                        // Check stop tokens
                        if let stopTokens = config.stopTokens, stopTokens.contains(nextToken) {
                            break
                        }
                        
                        // Check EOS (handle multiple EOS tokens)
                        if tokenizer.eosTokenIds.contains(nextToken) {
                            break
                        }
                        
                        // Add token
                        generatedTokens.append(nextToken)
                        
                        // Decode using streaming decoder (handles multi-byte UTF-8 correctly)
                        let tokenText = streamDecoder.consume(tokenId: nextToken)
                        if !tokenText.isEmpty {
                            continuation.yield(tokenText)
                        }
                        
                        // Forward pass for next token
                        lastLogits = try model.forward(tokenId: nextToken, position: position)
                        position += 1
                    }
                    
                    continuation.finish()
                    
                } catch {
                    continuation.finish()
                }
            }
        }
    }
    
    // ─────────────────────────────────────────────────────────────
    // Complete Generate - Returns full text
    // ─────────────────────────────────────────────────────────────
    
    public func generateComplete(
        prompt: String,
        config: GenerationConfig = .defaultConfig
    ) throws -> String {
        print("[GEN COMPLETE] Starting generation for prompt: '\(prompt)'")
        fflush(stdout)
        
        // Encode prompt
        print("[GEN COMPLETE] Encoding prompt...")
        fflush(stdout)
        let promptTokens = tokenizer.encode(text: prompt)
        print("[GEN COMPLETE] Encoded to \(promptTokens.count) tokens: \(promptTokens)")
        fflush(stdout)
        
        // Pre-fill KV cache with prompt tokens
        print("[GEN COMPLETE] Starting forward pass for prompt tokens...")
        fflush(stdout)
        var lastLogits: [Float] = []
        for (position, tokenId) in promptTokens.enumerated() {
            print("[GEN COMPLETE] Forward pass for token \(tokenId) at position \(position)")
            fflush(stdout)
            lastLogits = try model.forward(tokenId: tokenId, position: position)
            print("[GEN COMPLETE] Forward pass completed, logits count: \(lastLogits.count)")
            fflush(stdout)
        }
        print("[GEN COMPLETE] All prompt tokens processed")
        fflush(stdout)
        
        // Generate tokens
        var generatedTokens: [Int] = []
        var position = promptTokens.count
        
        for _ in 0..<config.maxTokens {
            let nextToken = sampler.sample(
                logits: lastLogits,
                temperature: config.temperature,
                topK: config.topK,
                topP: config.topP
            )
            
            // Debug: print selected token
            if generatedTokens.count < 5 {
                print("[DEBUG generateComplete] Token \(generatedTokens.count): ID=\(nextToken), raw='\(tokenizer.rawToken(for: nextToken) ?? "nil")'")
                fflush(stdout)
                // Print logits stats
                let maxLogit = lastLogits.max() ?? 0
                let minLogit = lastLogits.min() ?? 0
                let maxIdx = lastLogits.indices.filter { lastLogits[$0] == maxLogit }.first ?? -1
                print("[DEBUG] Logits: max=\(maxLogit) at idx=\(maxIdx), min=\(minLogit)")
                fflush(stdout)
            }
            
            if let stopTokens = config.stopTokens, stopTokens.contains(nextToken) {
                break
            }
            
            if tokenizer.eosTokenIds.contains(nextToken) {
                break
            }
            
            generatedTokens.append(nextToken)
            
            // Forward pass for next token
            lastLogits = try model.forward(tokenId: nextToken, position: position)
            position += 1
        }
        
        // Decode full response
        return tokenizer.decode(tokens: generatedTokens)
    }
    
    // ─────────────────────────────────────────────────────────────
    // Generate with Token IDs - Returns token sequence
    // ─────────────────────────────────────────────────────────────
    
    public func generateTokens(
        promptTokens: [Int],
        config: GenerationConfig = .defaultConfig
    ) throws -> [Int] {
        // Pre-fill KV cache with prompt tokens
        var lastLogits: [Float] = []
        for (position, tokenId) in promptTokens.enumerated() {
            lastLogits = try model.forward(tokenId: tokenId, position: position)
        }
        
        var generatedTokens: [Int] = []
        var position = promptTokens.count
        
        for _ in 0..<config.maxTokens {
            let nextToken = sampler.sample(
                logits: lastLogits,
                temperature: config.temperature,
                topK: config.topK,
                topP: config.topP
            )
            
            if let stopTokens = config.stopTokens, stopTokens.contains(nextToken) {
                break
            }
            
            if tokenizer.eosTokenIds.contains(nextToken) {
                break
            }
            
            generatedTokens.append(nextToken)
            
            // Forward pass for next token
            lastLogits = try model.forward(tokenId: nextToken, position: position)
            position += 1
        }
        
        return generatedTokens
    }
}