RWKVInferenceExample.swift

//
//  RWKVInferenceExample.swift
//  Squirrel
//
//  Example usage of RWKV inference for input method enhancement
//

import Foundation

/// Example class showing how to use RWKV inference in an input method
class RWKVInferenceExample {
    
    private var rwkv: RWKVInference?
    
    /// Initialize with model and vocab paths
    init() {
        // Get the application bundle path
        guard let bundlePath = Bundle.main.resourcePath else {
            NSLog("[RWKVExample] Failed to get bundle path")
            return
        }
        
        // Construct paths to model and vocab files
        let modelPath = bundlePath + "/rwkv-ime_L6D512_260107.gguf"
        let vocabPath = bundlePath + "/rwkv_vocab_ime.txt"
        
        // Create inference engine
        rwkv = RWKVInference(modelPath: modelPath, vocabPath: vocabPath)
        
        // Load model
        if let engine = rwkv {
            let success = engine.loadModel()
            if success {
                NSLog("[RWKVExample] RWKV model loaded successfully")
                
                // Set optimal parameters for input method prediction
                engine.setSamplerParams(
                    temperature: 0.8,    // Slightly creative
                    topK: 10,            // Consider top 10 tokens
                    topP: 0.9            // Nucleus sampling
                )
                
                engine.setPenaltyParams(
                    presencePenalty: 0.1,     // Slight penalty for repeating
                    frequencyPenalty: 0.1,    // Slight penalty for frequent tokens
                    penaltyDecay: 0.996       // Decay penalty over time
                )
            } else {
                NSLog("[RWKVExample] Failed to load RWKV model")
                rwkv = nil
            }
        }
    }
    
    deinit {
        rwkv?.unloadModel()
    }
    
    /// Get text predictions for input method
    /// - Parameter input: User's current input text
    /// - Returns: Array of prediction candidates
    func getPredictions(for input: String) -> [String] {
        guard let engine = rwkv else {
            NSLog("[RWKVExample] RWKV engine not initialized")
            return []
        }
        
        // For IME, we typically want top-K single token predictions
        let candidates = engine.getTopKPredictions(prompt: input, topK: 5)
        
        return candidates
    }
    
    /// Generate longer completion (for sentence prediction)
    /// - Parameters:
    ///   - input: User's current input text
    ///   - maxLength: Maximum length of completion
    /// - Returns: Generated completion text
    func generateCompletion(for input: String, maxLength: Int32 = 20) -> String? {
        guard let engine = rwkv else {
            NSLog("[RWKVExample] RWKV engine not initialized")
            return nil
        }
        
        return engine.generateCompletion(prompt: input, maxTokens: maxLength)
    }
    
    /// Clear the model state (useful when starting a new sentence)
    func resetState() {
        rwkv?.clearState()
        NSLog("[RWKVExample] Model state reset")
    }
    
    // MARK: - Example Usage
    
    /// Run a simple test to verify the inference works
    static func runTest() {
        NSLog("[RWKVExample] Starting RWKV inference test...")
        
        let example = RWKVInferenceExample()
        
        // Test 1: Get single token predictions
        NSLog("[RWKVExample] Test 1: Top-K predictions")
        let testInput1 = "你好"
        let predictions = example.getPredictions(for: testInput1)
        NSLog("[RWKVExample] Input: '\(testInput1)'")
        NSLog("[RWKVExample] Predictions: \(predictions)")
        
        // Test 2: Generate longer completion
        NSLog("[RWKVExample] Test 2: Text completion")
        let testInput2 = "今天天气"
        if let completion = example.generateCompletion(for: testInput2, maxLength: 10) {
            NSLog("[RWKVExample] Input: '\(testInput2)'")
            NSLog("[RWKVExample] Completion: \(completion)")
        }
        
        // Test 3: Reset state and try again
        NSLog("[RWKVExample] Test 3: Reset and predict")
        example.resetState()
        let predictions2 = example.getPredictions(for: "早上好")
        NSLog("[RWKVExample] Predictions after reset: \(predictions2)")
        
        NSLog("[RWKVExample] Test completed")
    }
}

// MARK: - Integration with SquirrelInputController

extension SquirrelInputController {
    
    /// Initialize RWKV inference (call this when input controller is created)
    func initializeRWKVInference() {
        // This can be called in viewDidLoad or init
        // Uncomment to enable RWKV inference
        // RWKVInferenceExample.runTest()
    }
}