v2: Initial clean branch with unit tests + CI/CD pipeline

- Started from ac75faa (initial E4B-MarkBase integration) - Kept Sources/ (all engine code) + Package.swift + .gitignore - Removed all ad-hoc tests, documentation, scripts, Python files - Added Tests/00_Unit/ (MathTest, TokenizerTest, SamplerTest) - Added .gitea/workflows/ci.yaml (build + unit tests + lint) - Added Scripts/check_resources.sh (memory-aware test runner) - Added Tests/Manifest.json (resource requirements for all tests) - Focus: 4-bit quantized models only
2026-07-05 13:29:25 +08:00
commit 8a66b9086a
90 changed files with 22252 additions and 0 deletions
@@ -0,0 +1,157 @@
+import Foundation
+
+// ─────────────────────────────────────────────────────────────
+// Sampler - Token sampling strategies
+// ─────────────────────────────────────────────────────────────
+
+public final class Sampler: @unchecked Sendable {
+    
+    public init() {}
+    
+    // ─────────────────────────────────────────────────────────────
+    // Sample - Main sampling function
+    // ─────────────────────────────────────────────────────────────
+    
+    public func sample(
+        logits: [Float],
+        temperature: Float = 1.0,
+        topK: Int? = nil,
+        topP: Float? = nil,
+        filterUnusedTokens: Bool = true,
+        unusedTokenRange: Range<Int> = 258000..<259000
+    ) -> Int {
+        var filteredLogits = logits
+        
+        // Filter out unused tokens if requested
+        if filterUnusedTokens {
+            for i in unusedTokenRange {
+                if i < filteredLogits.count {
+                    filteredLogits[i] = -Float.infinity
+                }
+            }
+        }
+        
+        // Handle temperature=0.0 (greedy sampling)
+        if temperature == 0.0 {
+            return greedySample(logits: filteredLogits)
+        }
+        
+        // Apply temperature
+        var scaledLogits = filteredLogits.map { $0 / temperature }
+        
+        // Apply Top-k
+        if let k = topK {
+            scaledLogits = applyTopK(logits: scaledLogits, k: k)
+        }
+        
+        // Apply Top-p (nucleus)
+        if let p = topP {
+            scaledLogits = applyTopP(logits: scaledLogits, p: p)
+        }
+        
+        // Convert to probabilities
+        let probs = softmax(logits: scaledLogits)
+        
+        // Random sample
+        return randomSample(probs: probs)
+    }
+    
+    // ─────────────────────────────────────────────────────────────
+    // Greedy Sample - Maximum probability
+    // ─────────────────────────────────────────────────────────────
+    
+    public func greedySample(logits: [Float]) -> Int {
+        var maxValue = logits[0]
+        var maxIndex = 0
+        
+        for i in 1..<logits.count {
+            if logits[i] > maxValue {
+                maxValue = logits[i]
+                maxIndex = i
+            }
+        }
+        
+        return maxIndex
+    }
+    
+    // ─────────────────────────────────────────────────────────────
+    // Top-k Filtering - Keep top k tokens
+    // ─────────────────────────────────────────────────────────────
+    
+    private func applyTopK(logits: [Float], k: Int) -> [Float] {
+        // Find threshold for top-k
+        let sorted = logits.sorted(by: >)
+        let threshold = sorted[min(k - 1, sorted.count - 1)]
+        
+        // Filter logits
+        return logits.map { logit in
+            logit >= threshold ? logit : -Float.infinity
+        }
+    }
+    
+    // ─────────────────────────────────────────────────────────────
+    // Top-p Filtering - Nucleus sampling
+    // ─────────────────────────────────────────────────────────────
+    
+    private func applyTopP(logits: [Float], p: Float) -> [Float] {
+        // Convert to probabilities
+        let probs = softmax(logits: logits)
+        
+        // Sort by probability
+        let sortedIndices = probs.indices.sorted { probs[$0] > probs[$1] }
+        
+        // Find cutoff
+        var cumulativeProb: Float = 0.0
+        var cutoffIndex = 0
+        
+        for idx in sortedIndices {
+            cumulativeProb += probs[idx]
+            if cumulativeProb >= p {
+                cutoffIndex = idx
+                break
+            }
+        }
+        
+        // Filter logits
+        return logits.indices.map { i in
+            probs[i] >= probs[cutoffIndex] ? logits[i] : -Float.infinity
+        }
+    }
+    
+    // ─────────────────────────────────────────────────────────────
+    // Softmax - Convert logits to probabilities
+    // ─────────────────────────────────────────────────────────────
+    
+    private func softmax(logits: [Float]) -> [Float] {
+        // Find max for numerical stability
+        let maxLogit = logits.max() ?? 0
+        
+        // Compute exp
+        let exps = logits.map { exp($0 - maxLogit) }
+        
+        // Normalize
+        let sum = exps.reduce(0, +)
+        return exps.map { $0 / sum }
+    }
+    
+    // ─────────────────────────────────────────────────────────────
+    // Random Sample - Sample from probability distribution
+    // ─────────────────────────────────────────────────────────────
+    
+    private func randomSample(probs: [Float]) -> Int {
+        // Generate random number
+        let rand = Float.random(in: 0..<1)
+        
+        // Find corresponding token
+        var cumulative: Float = 0.0
+        for i in probs.indices {
+            cumulative += probs[i]
+            if rand < cumulative {
+                return i
+            }
+        }
+        
+        // Fallback to last token
+        return probs.count - 1
+    }
+}