Code search demo: ask a codebase in plain English (EmbeddingGemma 2, Neural Engine)

Package.swift

@@ -54,6 +54,10 @@ let package = Package(
         .executableTarget(name: "BookmarkSortCheck", dependencies: ["BookmarkSort", "FluidUse"]),
         .executableTarget(name: "BookmarkSortDemo", dependencies: ["BookmarkSort"]),
         .executableTarget(name: "AudioSearchCheck", dependencies: ["FluidUse"]),
+        .target(name: "CodeSearch"),
+        .executableTarget(name: "CodeSearchCheck", dependencies: ["CodeSearch", "FluidUse"]),
+        .executableTarget(
+            name: "CodeSearchDemo", dependencies: ["CodeSearch", "FluidUse"], exclude: ["README.md", "demo.sh"]),
         .executableTarget(name: "AudioSearchDemo", dependencies: ["FluidUse"], exclude: ["README.md", "demo.sh"]),
         .executableTarget(
             name: "TopicSortDemo", dependencies: ["BookmarkSort", "FluidUse"], exclude: ["README.md", "demo.sh"],
@@ -81,6 +85,7 @@ let package = Package(
         ),
         .testTarget(name: "LayaTetrisTests", dependencies: ["LayaTetris"]),
         .testTarget(name: "BookmarkSortTests", dependencies: ["BookmarkSort"]),
+        .testTarget(name: "CodeSearchTests", dependencies: ["CodeSearch"]),
         .testTarget(name: "LayaTetrisDemoTests", dependencies: ["LayaTetrisDemo", "LayaTetris"]),
         .testTarget(name: "Game2048Tests", dependencies: ["Game2048"]),
         .testTarget(name: "GLiClass2048DemoTests", dependencies: ["GLiClass2048Demo", "Game2048"]),

Sources/CodeSearch/CodeChunker.swift

@@ -0,0 +1,102 @@
+import Foundation
+
+/// One searchable piece of code: a function, initializer, or type declaration with its doc comment.
+public struct CodeChunk: Sendable, Hashable {
+    /// Path relative to the repository root.
+    public let path: String
+    /// 1-based line of the declaration.
+    public let line: Int
+    /// `Type.member`, or the type's own name.
+    public let name: String
+    public let kind: String
+    /// Doc comment and code, up to `CodeChunker.maxLines` lines.
+    public let code: String
+
+    /// The document text EmbeddingGemma 2 was trained on for code retrieval: `title: <file> <symbol> | text: <code>`.
+    public var document: String { "title: \((path as NSString).lastPathComponent) \(name) | text: \(code)" }
+}
+
+/// Splits Swift sources into declaration chunks with a line scanner (no compiler): each `func`, `init`,
+/// `subscript`, or type declaration starts a chunk that runs, doc comment first, to the next declaration.
+public enum CodeChunker {
+    public static let maxLines = 40
+
+    private static let declaration = try! NSRegularExpression(
+        pattern:
+            #"^(\s*)(?:@[A-Za-z_]+(?:\([^)]*\))?\s+)*(?:(?:public|private|internal|fileprivate|open|static|final|override|mutating|nonmutating|nonisolated|convenience|required|indirect|package|class|distributed)\s+|(?:private|fileprivate|internal|public)\(set\)\s+)*(func|init|deinit|subscript|struct|class|enum|actor|protocol|extension)\b\s*([A-Za-z_][A-Za-z0-9_.]*)?"#
+    )
+    private static let typeKinds: Set<String> = ["struct", "class", "enum", "actor", "protocol", "extension"]
+    private static let skippedDirectories: Set<String> = [".build", ".git", ".swiftpm", "Packages", "DerivedData"]
+
+    /// Chunks of every `.swift` file under `repository`, in path order.
+    public static func chunks(repository: URL) -> [CodeChunk] {
+        let root = repository.standardizedFileURL
+        guard
+            let walker = FileManager.default.enumerator(
+                at: root, includingPropertiesForKeys: [.isDirectoryKey], options: [.skipsHiddenFiles])
+        else { return [] }
+        var files: [URL] = []
+        for case let url as URL in walker {
+            if skippedDirectories.contains(url.lastPathComponent) {
+                walker.skipDescendants()
+                continue
+            }
+            if url.pathExtension == "swift" { files.append(url) }
+        }
+        let prefix = root.path.hasSuffix("/") ? root.path : root.path + "/"
+        return files.sorted { $0.path < $1.path }.flatMap { file -> [CodeChunk] in
+            guard let source = try? String(contentsOf: file, encoding: .utf8) else { return [] }
+            let path = file.standardizedFileURL.path.replacingOccurrences(of: prefix, with: "")
+            return chunks(source: source, path: path)
+        }
+    }
+
+    /// Chunks of one file's source.
+    public static func chunks(source: String, path: String) -> [CodeChunk] {
+        let lines = source.components(separatedBy: "\n")
+        // (line index, indent, kind, name)
+        var declarations: [(index: Int, indent: Int, kind: String, name: String)] = []
+        for (index, line) in lines.enumerated() {
+            let range = NSRange(line.startIndex..., in: line)
+            guard let match = declaration.firstMatch(in: line, range: range),
+                let kindRange = Range(match.range(at: 2), in: line)
+            else { continue }
+            let kind = String(line[kindRange])
+            // `class func` / `class var` are members, not types.
+            if kind == "class",
+                line.range(of: #"\bclass\s+(func|var|let|subscript)\b"#, options: .regularExpression) != nil
+            {
+                continue
+            }
+            let indent = Range(match.range(at: 1), in: line).map { line[$0].count } ?? 0
+            var name = Range(match.range(at: 3), in: line).map { String(line[$0]) } ?? kind
+            if kind == "init" || kind == "deinit" || kind == "subscript" { name = kind }
+            declarations.append((index, indent, kind, name))
+        }
+        var result: [CodeChunk] = []
+        var types: [(indent: Int, name: String)] = []
+        for (position, item) in declarations.enumerated() {
+            while let last = types.last, last.indent >= item.indent { types.removeLast() }
+            let owner = types.last?.name
+            if typeKinds.contains(item.kind) { types.append((item.indent, item.name)) }
+            // Doc comment and attributes directly above the declaration.
+            var first = item.index
+            while first > 0 {
+                let above = lines[first - 1].trimmingCharacters(in: .whitespaces)
+                guard above.hasPrefix("///") || above.hasPrefix("@") || above.hasPrefix("*") || above.hasPrefix("/**")
+                else { break }
+                first -= 1
+            }
+            let next = position + 1 < declarations.count ? declarations[position + 1].index : lines.count
+            var last = max(item.index, next - 1)
+            // Leave the next declaration's doc comment to it.
+            while last > item.index, lines[last].trimmingCharacters(in: .whitespaces).hasPrefix("///") { last -= 1 }
+            let body = lines[first...min(last, first + maxLines - 1)]
+            let code = body.map { $0.trimmingCharacters(in: .whitespaces) }.filter { !$0.isEmpty }
+                .joined(separator: "\n")
+            let name = typeKinds.contains(item.kind) || owner == nil ? item.name : "\(owner!).\(item.name)"
+            result.append(CodeChunk(path: path, line: item.index + 1, name: name, kind: item.kind, code: code))
+        }
+        return result
+    }
+}

Sources/CodeSearchCheck/main.swift

@@ -0,0 +1,115 @@
+import CodeSearch
+import Foundation
+import FluidUse
+
+/// Indexes a repository's Swift declarations with EmbeddingGemma 2 and searches them in plain English.
+///
+///     swift run -c release CodeSearchCheck [--repo=path] [--tokens=64,128,256] [--eval] [--query="…"]…
+///
+/// `--eval` scores FluidAudio questions whose answer is known (top-1 / top-5 by file and name).
+@main
+struct CodeSearchCheck {
+    struct Question {
+        let text: String
+        let file: String
+        let name: String?
+    }
+
+    /// Questions about FluidAudio, worded without the function's own name.
+    static let questions = [
+        Question(
+            text: "resample an audio file to 16 kHz mono", file: "AudioConverter.swift", name: "resampleAudioFile"),
+        Question(
+            text: "download every file of a Hugging Face model repository", file: "DownloadUtils.swift",
+            name: "downloadRepo"),
+        Question(text: "variational Bayes clustering of speaker embeddings", file: "VBxClustering.swift", name: nil),
+        Question(text: "softmax over an array of floats", file: "VDSPOperations.swift", name: "softmax"),
+        Question(
+            text: "merge the transcripts of overlapping audio chunks", file: "ChunkProcessor.swift", name: "mergeChunks"
+        ),
+        Question(
+            text: "cosine distance between two speaker embeddings", file: "SpeakerOperations.swift",
+            name: "cosineDistance"),
+        Question(
+            text: "edit distance between two sequences", file: "LevenshteinDistance.swift", name: "levenshteinDistance"),
+        Question(text: "write float samples to a WAV file", file: "AudioConverter.swift", name: "AudioWAV"),
+        Question(text: "agglomerative hierarchical clustering", file: "AHCClustering.swift", name: nil),
+        Question(text: "convert a word into phonemes", file: "G2PModel.swift", name: "phonemize"),
+        Question(
+            text: "text to speech with the Kokoro model on the Neural Engine", file: "KokoroAneManager.swift",
+            name: "synthesize"),
+        Question(
+            text: "greedy transducer decoding that also returns token timestamps", file: "AsrManager.swift",
+            name: "tdtDecodeWithTimings"),
+        Question(text: "k-means clustering", file: "KMeansClustering.swift", name: nil),
+        Question(
+            text: "stretch a weight curve to a new length with interpolation", file: "WeightInterpolation.swift",
+            name: "resample"),
+        Question(
+            text: "load the reference recording used to clone a voice", file: "PocketTtsVoiceCloner.swift",
+            name: "loadAudio"),
+        Question(
+            text: "split a phoneme string into tokens for SSML", file: "SSMLProcessor.swift", name: "tokenizePhonemes"),
+    ]
+
+    static func main() async throws {
+        let arguments = CommandLine.arguments.dropFirst()
+        func value(_ name: String) -> String? {
+            arguments.first { $0.hasPrefix("--\(name)=") }.map { String($0.dropFirst(name.count + 3)) }
+        }
+        let repo = ((value("repo") ?? "~/Documents/FluidAudio") as NSString).expandingTildeInPath
+        let caps = (value("tokens") ?? "128").split(separator: ",").compactMap { Int($0) }
+        let queries = arguments.filter { $0.hasPrefix("--query=") }.map { String($0.dropFirst(8)) }
+        var start = DispatchTime.now().uptimeNanoseconds
+        let chunks = CodeChunker.chunks(repository: URL(fileURLWithPath: repo))
+        print(
+            String(
+                format: "%d chunks from %d files in %.2f s", chunks.count, Set(chunks.map(\.path)).count,
+                seconds(since: start)))
+        let manager = try await EmbeddingGemma2Manager.loadDefault()
+        _ = try await manager.embed(Array(repeating: "warm up", count: 16), prompt: .none)
+        for cap in caps {
+            start = DispatchTime.now().uptimeNanoseconds
+            var vectors: [[Float]] = []
+            for batch in stride(from: 0, to: chunks.count, by: 256) {
+                let slice = chunks[batch..<min(batch + 256, chunks.count)]
+                vectors += try await manager.embed(slice.map(\.document), prompt: .none, maxTokens: cap)
+            }
+            let indexSeconds = seconds(since: start)
+            print(
+                String(
+                    format: "cap %d tokens: indexed %d chunks in %.1f s (%.0f chunks/s)", cap, chunks.count,
+                    indexSeconds, Double(chunks.count) / indexSeconds))
+            func top(_ query: String, _ k: Int) async throws -> [CodeChunk] {
+                let q = try await manager.embed(query, prompt: .codeRetrieval)
+                let scored = vectors.indices.map { ($0, zip(q, vectors[$0]).reduce(Float(0)) { $0 + $1.0 * $1.1 }) }
+                return scored.sorted { $0.1 > $1.1 }.prefix(k).map { chunks[$0.0] }
+            }
+            if arguments.contains("--eval") {
+                var hit1 = 0
+                var hit5 = 0
+                for question in questions {
+                    let found = try await top(question.text, 5)
+                    let match = { (chunk: CodeChunk) in
+                        chunk.path.hasSuffix(question.file)
+                            && (question.name.map { chunk.name.hasSuffix($0) || chunk.name.contains($0) } ?? true)
+                    }
+                    if let first = found.first, match(first) { hit1 += 1 }
+                    if found.contains(where: match) { hit5 += 1 }
+                    if !arguments.contains("--quiet") {
+                        print(
+                            "  \(found.first.map(match) == true ? "✓" : found.contains(where: match) ? "~" : "✗") \(question.text) → \(found.first.map { "\($0.name) (\(($0.path as NSString).lastPathComponent):\($0.line))" } ?? "-")"
+                        )
+                    }
+                }
+                print("  top-1 \(hit1)/\(questions.count), top-5 \(hit5)/\(questions.count)")
+            }
+            for query in queries {
+                print("\(query):")
+                for chunk in try await top(query, 3) { print("   \(chunk.name)  \(chunk.path):\(chunk.line)") }
+            }
+        }
+    }
+
+    static func seconds(since start: UInt64) -> Double { Double(DispatchTime.now().uptimeNanoseconds - start) / 1e9 }
+}

Sources/CodeSearchDemo/CodeSearchDemoApp.swift

@@ -0,0 +1,24 @@
+import AppKit
+import SwiftUI
+
+@main
+struct CodeSearchDemoApp: App {
+    @StateObject private var model = CodeSearchModel()
+
+    init() {
+        setvbuf(stdout, nil, _IOLBF, 0)
+        // Bare SwiftPM executables start as background processes; make this one a regular windowed app.
+        NSApplication.shared.setActivationPolicy(.regular)
+        NSApplication.shared.activate(ignoringOtherApps: true)
+    }
+
+    var body: some Scene {
+        WindowGroup("Code search — on-device") {
+            ContentView()
+                .environmentObject(model)
+                .frame(minWidth: 1000, minHeight: 700)
+                .onAppear { model.start() }
+        }
+        .defaultSize(width: 1200, height: 950)
+    }
+}