From 21870a18e2811abf825cb01299ad2e0556d606a5 Mon Sep 17 00:00:00 2001 From: shimon <63890274+shimon-ai-dev@users.noreply.github.com> Date: Sat, 17 Oct 2026 13:04:41 +0800 Subject: [PATCH] Fix Han body matching and tokenizer-budget text tails Tokenize Han bodies into contiguous FTS phrases and preserve original bounded snippets. Split complete prompts using actual runtime token counts, preserve source locations and continuation, and reject oversized embedding requests instead of silently truncating them. Schema 5 to 6 rebuilds stored bodies without rereading source files or removing vectors. Previously truncated vectors need explicit reindex; older binaries cannot open schema 6. Add deterministic, worker-contract and actual-model regressions. Preserve upstream identity, filename-prefix behavior and dependency lockfile. --- README.md | 11 +- Sources/DigUpKit/CodeExtractor.swift | 7 +- Sources/DigUpKit/Extractors.swift | 4 +- Sources/DigUpKit/IndexStore.swift | 28 +- Sources/DigUpKit/Indexer.swift | 94 +++++-- Sources/DigUpKit/SearchText.swift | 54 +++- Sources/DigUpKit/Searcher.swift | 23 +- Sources/DigUpKit/TextBudget.swift | 58 ++++ Sources/DigUpKit/Worker.swift | 34 ++- Sources/LlamaRuntime/LlamaEmbedder.swift | 17 +- Sources/digup/IndexService.swift | 2 + Sources/digup/WorkerService.swift | 4 + Tests/DigUpKitTests/ChineseSearchTests.swift | 254 ++++++++++++++++++ Tests/DigUpKitTests/CodeTests.swift | 69 ++++- Tests/DigUpKitTests/FakeEmbedder.swift | 2 + Tests/DigUpKitTests/TextBudgetTests.swift | 36 +++ Tests/DigUpKitTests/WorkerBudgetTests.swift | 56 ++++ .../NativeTextBudgetTests.swift | 36 +++ docs/han-body-search.md | 47 ++++ worker/digup_worker.py | 14 +- worker/test_text_budget.py | 57 ++++ 21 files changed, 839 insertions(+), 68 deletions(-) create mode 100644 Sources/DigUpKit/TextBudget.swift create mode 100644 Tests/DigUpKitTests/ChineseSearchTests.swift create mode 100644 Tests/DigUpKitTests/TextBudgetTests.swift create mode 100644 Tests/DigUpKitTests/WorkerBudgetTests.swift create mode 100644 Tests/LlamaRuntimeTests/NativeTextBudgetTests.swift create mode 100644 docs/han-body-search.md create mode 100644 worker/test_text_budget.py diff --git a/README.md b/README.md index 219c8043..05db03ee 100644 --- a/README.md +++ b/README.md @@ -205,11 +205,12 @@ carry their own PDF, video and image readers. **Will it slow my computer down?** Only the first pass is heavy. After that it reads new files only. The model runs in a separate process that quits when it's done, and the search side only loads while you're searching. -**Which languages?** The model knows more than 100, and DigUp searches in all of them, across languages too. The evals -cover only English, Bengali and Arabic because those are the languages of the files DigUp was tested on, and Google -notes that the model isn't equally strong in every language. Two limits apply to exact words, not to meaning: the -text inside screenshots is read only on a Mac, by Apple's text recognition, which knows about 25 languages, and -languages written without spaces, like Chinese, Japanese and Thai, aren't split into words. +**Which languages?** The model knows more than 100, and DigUp searches in all of them, across languages too. The semantic +evals cover only English, Bengali and Arabic because those are the languages of the files DigUp was tested on, and +Google notes that the model isn't equally strong in every language. Exact-word search reads screenshot text only on +a Mac, using Apple's text recognition (about 25 languages). Han characters in text bodies match as contiguous +sequences; other text written without spaces, such as Japanese kana and Thai, is not split into words. See +[Han body search](docs/han-body-search.md) for index-upgrade details. **How do I remove it?** On a Mac, quit it, then delete DigUp.app and `~/Library/Application Support/DigUp`, which holds the model and the diff --git a/Sources/DigUpKit/CodeExtractor.swift b/Sources/DigUpKit/CodeExtractor.swift index 0dfc0010..80762a96 100644 --- a/Sources/DigUpKit/CodeExtractor.swift +++ b/Sources/DigUpKit/CodeExtractor.swift @@ -23,9 +23,6 @@ public enum CodeExtractor { let cells: Int? } - /// The model reads 2,048 tokens of an input at most; a single line longer than this is cut here. - static let longestPiece = 4000 - static func read(_ url: URL, options: IndexOptions) throws -> File { let data = try Data(contentsOf: url) let (repo, title) = place(of: url.path) @@ -117,8 +114,8 @@ public enum CodeExtractor { } var last = end while last > start + 1, blank(last - 1) { last -= 1 } // blank lines at its end don't count - var body = lines[start.. longestPiece { body = String(body.prefix(longestPiece)) } + // The indexer splits complete prompts using the runtime tokenizer. Never discard a long line here. + let body = lines[start.. [SQLValue] { - let normalized = SearchText.normalized(text) + let normalized = SearchText.indexed(text) return [.text(normalized), normalized == text ? .null : .text(text)] } private func addNameRow(file: Int64, path: String) throws { try db.run("INSERT INTO segments(file, kind, modality) VALUES(?, 'name', 'none')", [.int(file)]) - try db.run("INSERT INTO fts(rowid, name, body) VALUES(?, ?, ?)", - [.int(db.lastInsertID), .text(SearchText.normalized(Self.nameText(path))), - .text(SearchText.normalized(Self.folderText(path)))]) + try db.run("INSERT INTO fts(rowid, name, body, original) VALUES(?, ?, ?, ?)", + [.int(db.lastInsertID), .text(SearchText.normalized(Self.nameText(path)))] + + Self.keywordText(Self.folderText(path))) } private func deleteSegments(of file: Int64, keepingName: Bool) throws { diff --git a/Sources/DigUpKit/Indexer.swift b/Sources/DigUpKit/Indexer.swift index 8c5bb405..3a83197b 100644 --- a/Sources/DigUpKit/Indexer.swift +++ b/Sources/DigUpKit/Indexer.swift @@ -583,28 +583,41 @@ public final class Indexer { let wordPages = textPages.filter { SearchText.hasWords($0.text) } let imagePages = pages.filter { $0.imagePath != nil } let garbledPages = pages.filter { $0.garbled != nil } // rendered too: among `imagePages` - var textVectors: [Int: [Float]] = [:] // by page number + var textParts: [(page: Int, text: String)] = [] + var textVectors: [Int: [Float]] = [:] // by part, not page (a page can have several) var imageVectors: [Embedding] = [] + var garbledParts: [(page: Int, text: String)] = [] var garbledVectors: [[Float]] = [] if !wordPages.isEmpty || !imagePages.isEmpty { let worker = try fullWorker() - let vectors = try worker.embedTexts(wordPages.map { - Prompts.document(title: title, text: String($0.text.prefix(4000))) + for page in textPages { + let parts = SearchText.hasWords(page.text) + ? try worker.documentChunks(page.text, title: title) : [page.text] + textParts += parts.map { (page.number, $0) } + } + let worded = textParts.indices.filter { SearchText.hasWords(textParts[$0].text) } + let vectors = try worker.embedTexts(worded.map { + Prompts.document(title: title, text: textParts[$0].text) }) - textVectors = Dictionary(uniqueKeysWithValues: zip(wordPages.map(\.number), vectors)) + textVectors = Dictionary(uniqueKeysWithValues: zip(worded, vectors)) imageVectors = try worker.embedImages(imagePages.compactMap(\.imagePath), budget: options.pageBudget) - garbledVectors = try worker.embedTexts(garbledPages.map { - Prompts.document(title: title, text: String($0.garbled!.prefix(4000))) + for page in garbledPages { + garbledParts += try worker.documentChunks(page.garbled!, title: title).map { (page.number, $0) } + } + garbledVectors = try worker.embedTexts(garbledParts.map { + Prompts.document(title: title, text: $0.text) }) + } else { + textParts = textPages.map { ($0.number, $0.text) } } try store.db.transaction { // Meaning only: the garbled words are never shown or matched against what you type. - for (page, vector) in zip(garbledPages, garbledVectors) { - try store.addSegment(file: file.id, kind: .page, modality: .text, loc: Double(page.number), vector: vector) + for (part, vector) in zip(garbledParts, garbledVectors) { + try store.addSegment(file: file.id, kind: .page, modality: .text, loc: Double(part.page), vector: vector) } - for page in textPages { - try store.addSegment(file: file.id, kind: .page, modality: .text, loc: Double(page.number), - excerpt: excerpt(page.text), text: page.text, vector: textVectors[page.number]) + for (index, part) in textParts.enumerated() { + try store.addSegment(file: file.id, kind: .page, modality: .text, loc: Double(part.page), + excerpt: excerpt(part.text), text: part.text, vector: textVectors[index]) } for (page, embedding) in zip(imagePages, imageVectors) { try store.addSegment(file: file.id, kind: .pageImage, modality: .image, loc: Double(page.number), @@ -634,18 +647,24 @@ public final class Indexer { } let title = URL(fileURLWithPath: file.path).deletingPathExtension().lastPathComponent // Text that's mostly digits (a data dump saved as .txt) gets keyword search only. - let worded = chunks.indices.filter { SearchText.hasWords(chunks[$0]) } - let vectors = worded.isEmpty ? [] : try fullWorker().embedTexts(worded.map { - Prompts.document(title: title, text: chunks[$0]) + let worker = chunks.contains(where: SearchText.hasWords) ? try fullWorker() : nil + var parts: [(offset: Int, text: String)] = [] + for (offset, chunk) in chunks.enumerated() { + let split = SearchText.hasWords(chunk) ? try worker!.documentChunks(chunk, title: title) : [chunk] + parts += split.map { (offset, $0) } + } + let worded = parts.indices.filter { SearchText.hasWords(parts[$0].text) } + let vectors = worded.isEmpty ? [] : try worker!.embedTexts(worded.map { + Prompts.document(title: title, text: parts[$0].text) }) let vector = Dictionary(uniqueKeysWithValues: zip(worded, vectors)) let end = min(total, options.maxChunks) let unread = read.upperBound < end try store.db.transaction { - for (offset, chunk) in chunks.enumerated() { + for (index, part) in parts.enumerated() { try store.addSegment(file: file.id, kind: .chunk, modality: .text, - loc: Double(read.lowerBound + offset + 1), excerpt: excerpt(chunk), text: chunk, - vector: vector[offset]) + loc: Double(read.lowerBound + part.offset + 1), excerpt: excerpt(part.text), text: part.text, + vector: vector[index]) } var info = ["chunks": "\(end)", "chars": "\(characters)"] info["read_to"] = unread ? "\(read.upperBound)" : nil @@ -656,7 +675,7 @@ public final class Indexer { let span = later ? "chunks \(read.lowerBound + 1)–\(read.upperBound) of \(end)" : unread ? "\(chunks.count) of \(end) chunks" : "\(chunks.count) chunks" return (unread ? "partial" : "done", - span + (worded.count < chunks.count ? ", \(chunks.count - worded.count) for keywords only" : "")) + span + (worded.count < parts.count ? ", \(parts.count - worded.count) for keywords only" : "")) } /// A code file's stretches of lines (a notebook's of cells), each embedded with the file's path in its repo as its @@ -666,14 +685,39 @@ public final class Indexer { try store.finish(file.id, state: "skipped", detail: "no code") return ("skipped", "no code") } - // Lines that are mostly numbers and symbols (a table of constants) get keyword search only. - let worded = code.pieces.indices.filter { SearchText.hasWords(code.pieces[$0].text) } - let vectors = worded.isEmpty ? [] : try fullWorker().embedTexts(worded.map { - Prompts.document(title: code.title, text: code.pieces[$0].text) + // Keep the extractor's coarse code boundaries, then apply the complete + // prompt's real token budget. Preserve source offsets through this split. + let worker = code.pieces.contains { SearchText.hasWords($0.text) } ? try fullWorker() : nil + var pieces: [CodeExtractor.Piece] = [] + for piece in code.pieces { + guard SearchText.hasWords(piece.text), let worker else { + pieces.append(piece) + continue + } + let ranges = try TextBudget.ranges(piece.text, title: code.title, limit: worker.textTokenLimit, + count: worker.textTokenCount) + var countedThrough = piece.text.startIndex, lineOffset = 0 + for range in ranges { + let text = String(piece.text[range]) + var first = piece.first, last = piece.last + if code.cells == nil { + // Ranges advance even with overlap; count each source prefix only once. + lineOffset += piece.text[countedThrough.. (String, String?) { diff --git a/Sources/DigUpKit/SearchText.swift b/Sources/DigUpKit/SearchText.swift index 35f06bad..8b521b22 100644 --- a/Sources/DigUpKit/SearchText.swift +++ b/Sources/DigUpKit/SearchText.swift @@ -2,6 +2,36 @@ import Foundation /// Text the way keyword search keeps it and matches it, so that a word is found however it was typed or extracted. public enum SearchText { + /// Body text and body query phrases only; filenames retain their unspaced representation. Han characters + /// become adjacent tokens, so "租房" matches inside "租房合同" without losing two-character words. + public static func indexed(_ text: String) -> String { + var result = "" + let characters = Array(normalized(text)) + for (index, character) in characters.enumerated() { + if isHan(character) { result += " " + String(character) + " " } + else { + result.append(character) + let boundary = !character.unicodeScalars.contains { CharacterSet.alphanumerics.contains($0) } + let besideHan = (index > 0 && isHan(characters[index - 1])) + || (index + 1 < characters.count && isHan(characters[index + 1])) + // Co is indexed but is never a query word. Source separators must + // not disappear between adjacent Han tokens. + if boundary && besideHan { result += " \u{E000} " } + } + } + return result.replacingOccurrences(of: #"[ \t]+"#, with: " ", options: .regularExpression) + .trimmingCharacters(in: .whitespaces) + } + + public static func isHan(_ character: Character) -> Bool { + character.unicodeScalars.contains { + switch $0.value { + case 0x3400...0x4DBF, 0x4E00...0x9FFF, 0xF900...0xFAFF, 0x20000...0x2FFFF, 0x30000...0x323AF: true + default: false + } + } + } + /// NFKC (Arabic presentation forms and ligatures become plain letters, full-width becomes ASCII, and Bengali's /// two spellings of "ো" become one), invisible formatting characters dropped (Bengali's zero-width joiners, /// direction marks), and Arabic the way people type it: no tatweel, no short-vowel marks, one alef. @@ -59,13 +89,16 @@ public enum SearchText { guard hit.location != NSNotFound else { break } searchRange = NSRange(location: hit.upperBound, length: haystack.length - hit.upperBound) guard var range = folded.original(hit) else { continue } - if range.lowerBound > text.startIndex, isWordCharacter(text[text.index(before: range.lowerBound)]) { + if word.first.map(isHan) != true, + range.lowerBound > text.startIndex, isWordCharacter(text[text.index(before: range.lowerBound)]), + !isHan(text[text.index(before: range.lowerBound)]) { continue } - if range.upperBound < text.endIndex, isWordCharacter(text[range.upperBound]) { + if word.last.map(isHan) != true, + range.upperBound < text.endIndex, isWordCharacter(text[range.upperBound]), !isHan(text[range.upperBound]) { guard prefix else { continue } var end = range.upperBound - while end < text.endIndex, isWordCharacter(text[end]) { end = text.index(after: end) } + while end < text.endIndex, isWordCharacter(text[end]), !isHan(text[end]) { end = text.index(after: end) } range = range.lowerBound.. String? { - let marks = terms.flatMap { matches(of: $0, in: text) }.sorted { $0.lowerBound < $1.lowerBound } + let marks = terms.flatMap { matches(of: $0, in: text) } + .sorted { $0.lowerBound < $1.lowerBound } guard let first = marks.first else { return nil } // About `words` words, starting a few before the first match. - var start = first.lowerBound, before = 0 - while start > text.startIndex, before < words / 3 { + let han = text.contains(where: isHan) + var start = first.lowerBound, before = 0, characters = 0 + while start > text.startIndex, before < words / 3, !han || characters < words * 3 { start = text.index(before: start) + characters += 1 if text[start].isWhitespace, start < first.lowerBound, !text[text.index(after: start)].isWhitespace { before += 1 } } if start > text.startIndex || text[start].isWhitespace { start = text.index(after: start) } var end = first.upperBound, after = 0 - while end < text.endIndex, after < words - before { + characters = 0 + while end < text.endIndex, after < words - before, !han || characters < words * 6 { if text[end].isWhitespace { after += 1 } end = text.index(after: end) + characters += 1 } var out = start > text.startIndex ? "…" : "" var cursor = start @@ -148,7 +186,7 @@ public enum Landmark { guard page > 1, texts.count >= page else { return nil } let target = texts[page - 1] let before = texts.prefix(page - 1).map { SearchText.normalized($0).lowercased() } - let onPage = words.filter { $0.count >= 3 }.compactMap { word in + let onPage = words.filter { $0.count >= 3 || ($0.count >= 2 && $0.contains(where: SearchText.isHan)) }.compactMap { word in SearchText.matches(of: word, in: target).first.map { String(target[$0]) } } return onPage.first { word in !before.contains { $0.contains(SearchText.normalized(word).lowercased()) } } diff --git a/Sources/DigUpKit/Searcher.swift b/Sources/DigUpKit/Searcher.swift index 4d7b6def..cc0325c7 100644 --- a/Sources/DigUpKit/Searcher.swift +++ b/Sources/DigUpKit/Searcher.swift @@ -143,8 +143,15 @@ public final class Searcher { // as typed stays whole, so "cat" doesn't bring in "category"; so does a word of two letters, the start of too // many words to stand for one. if identifier == nil { - for index in terms.indices where terms[index].count >= 3 && !patterns[index].hasSuffix("*") { - if try !isIndexed(patterns[index]) { patterns[index] += "*" } + for index in terms.indices where terms[index].count >= 3 + && !(prefixLast && index == terms.count - 1) && !SearchText.inflects(terms[index]) { + if try !isIndexed(patterns[index]) { + if terms[index].contains(where: SearchText.isHan) { + // Names are still whole Han tokens: retain their schema-5 missing-word prefix fallback. + // Body phrases must remain contiguous exact sequences of Han characters. + patterns[index] = "(name : \"\(terms[index])\"* OR body : \"\(SearchText.indexed(terms[index]))\")" + } else { patterns[index] += "*" } + } } } if !patterns.isEmpty { @@ -157,7 +164,8 @@ public final class Searcher { let file = row.int(0) guard words.count < 100, allowed(file), seen.insert(file).inserted else { return } // Text that normalizing changed is quoted as it's written (with its hamza, harakat, joiners). - let snippet = row.text(3).flatMap { SearchText.snippet($0, terms: terms) } ?? row.text(2) + // Stemming may match without a literal highlight. Never show the indexed Han separators. + let snippet = row.text(3).map { SearchText.snippet($0, terms: terms) ?? excerpt($0) } ?? row.text(2) words.append((file, row.int(1), snippet)) } if !words.isEmpty { @@ -228,7 +236,7 @@ public final class Searcher { // never shows page 2's words. A segment with no place and no text of its own (a screenshot's image // vector) can borrow them from its file's OCR text, which describes the same picture. var snippet = keyword.flatMap { $0.1.segment == segmentID ? $0.1.snippet : nil } - .flatMap { $0.contains("«") ? excerpt($0, limit: 400) : nil } + .flatMap { $0.contains("«") || segment != .name ? excerpt($0, limit: 400) : nil } if snippet == nil, storedExcerpt == nil, loc == nil, let match = keyword?.1, let text = match.snippet, text.contains("«") { var matchedKind: String? @@ -394,7 +402,12 @@ public final class Searcher { static func patterns(_ terms: [String], prefixLast: Bool) -> [String] { terms.enumerated().map { index, term in let prefix = (prefixLast && index == terms.count - 1) || SearchText.inflects(term) - return "\"\(term)\"" + (prefix ? "*" : "") + let suffix = prefix ? "*" : "" + if term.contains(where: SearchText.isHan) { + // Keep name matching (and the filename substring fallback) compatible with schema 5. + return "(name : \"\(term)\"\(suffix) OR body : \"\(SearchText.indexed(term))\"\(suffix))" + } + return "\"\(term)\"" + suffix } } diff --git a/Sources/DigUpKit/TextBudget.swift b/Sources/DigUpKit/TextBudget.swift new file mode 100644 index 00000000..ddc98664 --- /dev/null +++ b/Sources/DigUpKit/TextBudget.swift @@ -0,0 +1,58 @@ +import Foundation + +/// Splits documents against a complete prompt's tokenizer budget. Callers provide +/// the runtime's real token count; character counts never decide what fits. +public enum TextBudget { + /// A single file cannot fit; unlike a failed worker, this must not stop the indexing queue. + public struct Failure: Error, CustomStringConvertible { + public let description: String + } + public static func chunks(_ text: String, title: String?, limit: Int, + count: (String) throws -> Int) throws -> [String] { + try ranges(text, title: title, limit: limit, count: count).map { String(text[$0]) } + } + + /// Original offsets let code retain source lines even when repeated text is split. + static func ranges(_ text: String, title: String?, limit: Int, + count: (String) throws -> Int) throws -> [Range] { + guard !text.isEmpty else { return [] } + func fits(_ body: String) throws -> Bool { + try count(Prompts.document(title: title, text: body)) <= limit + } + if try fits(text) { return [text.startIndex..] = [] + while start < indices.count - 1 { + // Bound tokenizer work even for enormous pages. This bound is only + // a probe size: every emitted piece is checked with the tokenizer. + var low = start + 1, high = min(indices.count - 1, start + max(1800, limit * 4)) + var end = start + while low <= high { + let probe = low + (high - low) / 2 + if try fits(String(text[indices[start].. start else { + throw Failure(description: "document title and one character exceed the \(limit)-token budget") + } + // Prefer a word boundary when one is nearby; no-space Han stays overlapped. + if end < indices.count - 1, + let space = text[indices[start].. (end - start) / 2 { + let candidate = start + text.distance(from: indices[start], to: space) + if try fits(String(text[indices[start].. [[Float]] func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] func embedAudio(_ paths: [String]) throws -> [Embedding] + var textTokenLimit: Int { get } + func textTokenCount(_ text: String) throws -> Int /// Frees the model. Safe to call twice. func close() } +public extension Embedder { + var textTokenLimit: Int { info.textTokenLimit ?? 2048 } + /// Conservative byte bound for simple/test embedders. Production runtimes + /// override this with their tokenizer, including its special tokens. + func textTokenCount(_ text: String) throws -> Int { text.utf8.count + 2 } + func documentChunks(_ text: String, title: String?) throws -> [String] { + try TextBudget.chunks(text, title: title, limit: textTokenLimit, count: textTokenCount) + } +} + /// How to start a worker process that speaks the JSON-lines protocol (see `EmbeddingWorker`). public struct WorkerConfig: Sendable { public var executable: URL @@ -131,6 +145,7 @@ public final class EmbeddingWorker: Embedder { let vectors: String? let errors: [String?]? let info: WorkerInfo? + let counts: [Int]? } public init(_ config: WorkerConfig, textOnly: Bool) throws { @@ -162,6 +177,23 @@ public final class EmbeddingWorker: Embedder { return try send("text", texts, budget: nil).map { $0.vector ?? [] } } + public func textTokenCount(_ text: String) throws -> Int { + let id = nextID + nextID += 1 + let line = try JSONEncoder().encode(Request(id: id, op: "token_count", inputs: [text], budget: nil)) + do { + try input.fileHandleForWriting.write(contentsOf: line + [0x0A]) + } catch { + throw WorkerError.died("can't write to the worker\(logTail())") + } + let reply = try readReply() + guard reply.id == id, reply.ok == true, reply.counts?.count == 1, + let count = reply.counts?.first, count >= 0 else { + throw WorkerError.failed(reply.error ?? "worker does not report tokenizer counts") + } + return count + } + public func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] { guard !paths.isEmpty else { return [] } return try send("image", paths, budget: budget) diff --git a/Sources/LlamaRuntime/LlamaEmbedder.swift b/Sources/LlamaRuntime/LlamaEmbedder.swift index 952c78f1..0dcfcce0 100644 --- a/Sources/LlamaRuntime/LlamaEmbedder.swift +++ b/Sources/LlamaRuntime/LlamaEmbedder.swift @@ -58,7 +58,7 @@ public final class LlamaEmbedder: Embedder { throw WorkerError.failed("can't load \(models.text.path)") } // One input must fit one micro-batch: 2048 tokens (the reference's cut), which also holds a 1120-token image or - // a 750-token 30 s clip. Queries are short, so the query encoder reserves less and cuts at 512. + // a 750-token 30 s clip. Queries are short, so the query encoder reserves a smaller 512-token budget. let maxTokens = textOnly ? 512 : 2048 var params = llama_context_default_params() params.n_ctx = UInt32(maxTokens) @@ -104,7 +104,7 @@ public final class LlamaEmbedder: Embedder { self.maxTokens = maxTokens let known = Self.modelInfo info = WorkerInfo(model: known.model, revision: known.revision, runtime: known.runtime, dtype: known.dtype, - dim: embeddingSize, textOnly: textOnly, loadSeconds: nil) + dim: embeddingSize, textOnly: textOnly, loadSeconds: nil, textTokenLimit: maxTokens) // The first calls build the Metal pipelines: pay that now (the panel is opening). do { for _ in 0..<2 { _ = try embedTexts([Prompts.query("warm up")]) } @@ -114,7 +114,7 @@ public final class LlamaEmbedder: Embedder { } let seconds = (Date().timeIntervalSince(started) * 100).rounded() / 100 info = WorkerInfo(model: info.model, revision: info.revision, runtime: info.runtime, dtype: info.dtype, - dim: info.dim, textOnly: textOnly, loadSeconds: seconds) + dim: info.dim, textOnly: textOnly, loadSeconds: seconds, textTokenLimit: maxTokens) LlamaLog.shared.line("ready in \(seconds) s (\(textOnly ? "text-only" : "full"), llama.cpp \(Self.buildTag), " + "vectors \(Self.vectorVersion))") } @@ -136,6 +136,9 @@ public final class LlamaEmbedder: Embedder { var vectors: [[Float]] = [] for text in texts { let tokens = tokenize(text) + guard tokens.count <= maxTokens else { + throw WorkerError.failed("text has \(tokens.count) tokens; limit is \(maxTokens). Split it before embedding") + } if sequences == Self.maxSequences || self.tokens + tokens.count > maxTokens { vectors += try run() } begin() for token in tokens { try add(token: token) } @@ -144,6 +147,9 @@ public final class LlamaEmbedder: Embedder { return vectors } + public var textTokenLimit: Int { maxTokens } + public func textTokenCount(_ text: String) throws -> Int { tokenize(text).count } + public func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] { guard media != nil else { throw WorkerError.failed("this embedder was started text-only") } guard Self.budgets.contains(budget) else { throw WorkerError.failed("budget must be one of \(Self.budgets)") } @@ -247,7 +253,7 @@ public final class LlamaEmbedder: Embedder { // MARK: Text model - /// ` text `, cut to `maxTokens` the way transformers cuts (keeping both special tokens). + /// ` text `, never silently truncated. Embedding validates the budget. private func tokenize(_ text: String) -> [llama_token] { let utf8 = Array(text.utf8CString) var tokens = [llama_token](repeating: 0, count: utf8.count + 8) @@ -257,9 +263,6 @@ public final class LlamaEmbedder: Embedder { count = llama_tokenize(vocab, utf8, Int32(utf8.count - 1), &tokens, Int32(tokens.count), true, true) } tokens.removeLast(tokens.count - Int(max(count, 0))) - if tokens.count > maxTokens { - tokens = Array(tokens[0..<(maxTokens - 1)]) + [llama_vocab_eos(vocab)] - } return tokens } diff --git a/Sources/digup/IndexService.swift b/Sources/digup/IndexService.swift index 5c5a7102..94e18be5 100644 --- a/Sources/digup/IndexService.swift +++ b/Sources/digup/IndexService.swift @@ -368,6 +368,8 @@ final class SharedEmbedder: @unchecked Sendable { let model: any Embedder init(_ model: any Embedder) { self.model = model } var info: WorkerInfo { model.info } + var textTokenLimit: Int { model.textTokenLimit } + func textTokenCount(_ text: String) throws -> Int { try model.textTokenCount(text) } func embedTexts(_ texts: [String]) throws -> [[Float]] { try model.embedTexts(texts) } func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] { try model.embedImages(paths, budget: budget) diff --git a/Sources/digup/WorkerService.swift b/Sources/digup/WorkerService.swift index 47a7a8d1..10fe6594 100644 --- a/Sources/digup/WorkerService.swift +++ b/Sources/digup/WorkerService.swift @@ -17,6 +17,7 @@ func workerService(_ args: Arguments) -> Int32 { let infoFields: [String: Any] = [ "model": info.model, "revision": info.revision, "runtime": info.runtime, "dtype": info.dtype, "dim": info.dim, "text_only": info.textOnly, "load_seconds": info.loadSeconds ?? 0, + "text_token_limit": embedder.textTokenLimit, ] emit(["ready": true, "info": infoFields]) while let line = readLine() { @@ -36,6 +37,9 @@ func workerService(_ args: Arguments) -> Int32 { continue case "text": embeddings = try embedder.embedTexts(inputs).map { Embedding(vector: $0, error: nil) } + case "token_count": + emit(["id": id, "ok": true, "counts": try inputs.map(embedder.textTokenCount)]) + continue case "image": embeddings = try embedder.embedImages(inputs, budget: request["budget"] as? Int ?? 280) case "audio": diff --git a/Tests/DigUpKitTests/ChineseSearchTests.swift b/Tests/DigUpKitTests/ChineseSearchTests.swift new file mode 100644 index 00000000..8d25eb2d --- /dev/null +++ b/Tests/DigUpKitTests/ChineseSearchTests.swift @@ -0,0 +1,254 @@ +#if os(macOS) +import AppKit +#endif +import Foundation +import Testing +@testable import DigUpKit + +@Suite final class ChineseSearchTests { + let folder: URL = { + let caches = FileManager.default.urls(for: .cachesDirectory, in: .userDomainMask)[0] + let url = caches.appendingPathComponent("DigUpChineseTests-\(UUID().uuidString)") + try? FileManager.default.createDirectory(at: url, withIntermediateDirectories: true) + return url + }() + + deinit { try? FileManager.default.removeItem(at: folder) } + + @discardableResult + func add(_ store: IndexStore, path: String, body: String? = nil, inode: UInt64 = 1) throws -> Int64 { + let id = try store.insert(Candidate(path: path, kind: .doc, size: 1, mtime: 1, inode: inode, device: 1)) + if let body { try store.addSegment(file: id, kind: .chunk, modality: .text, loc: 1, text: body) } + return id + } + + @Test func chineseLexicalEvalOnPublicSyntheticCorpus() throws { + // Each case has its own index, so a keyword-only R@1 has one unambiguous answer. + let cases: [(query: String, name: String, body: String?)] = [ + ("租房", "notes.txt", "这是一份关于租房费用的公开合成记录"), + ("发票", "notes.txt", "本月发票金额已记录在公开合成表格中"), + ("预算", "notes.txt", "年度预算安排来自公开合成资料"), + ("北京", "notes.txt", "我们计划乘坐高铁前往北京游览"), + ("高铁", "notes.txt", "乘坐高铁前往南方旅行"), + ("苹果", "notes.txt", "桌上放着新鲜苹果和香蕉"), + ("𠮷野", "notes.txt", "公开合成记录中的𠮷野家店铺"), + ("年度租赁合同扫描备份", "年度租赁合同扫描备份.pdf", nil), + + ("payment error", "notes.txt", "A payment error appeared on the synthetic receipt."), + ] + var correct = 0 + for (index, item) in cases.enumerated() { + let store = try IndexStore(directory: folder.appendingPathComponent("eval-\(index)")) + let path = "/synthetic/\(item.name)" + try add(store, path: path, body: item.body) + let hits = try Searcher(store: store, loadVectors: false).search(item.query, queryVector: nil) + if hits.first?.path == path { correct += 1 } + #expect(hits.first?.path == path, "Keyword-only synthetic case: \(item.query)") + } + print("CHINESE_LEXICAL_EVAL keyword-only R@1: \(correct)/\(cases.count); no model inference") + } + + @Test func chinesePhrasesAreContiguousAndHighlightedAsWritten() throws { + let store = try IndexStore(directory: folder.appendingPathComponent("phrases")) + try add(store, path: "/synthetic/good.txt", body: "租房合同已经签署", inode: 1) + try add(store, path: "/synthetic/bad.txt", body: "租赁住房合同", inode: 2) + let hits = try Searcher(store: store, loadVectors: false).search("租房", queryVector: nil) + #expect(hits.map(\.path) == ["/synthetic/good.txt"]) + #expect(hits.first?.excerpt?.contains("«租房»合同") == true) + #expect(Landmark.find(page: 2, texts: ["公开资料", "这份租房合同已签署"], words: ["合同"]) == "合同") + #expect(SearchText.matches(of: "day", in: "Sunday day").map { String("Sunday day"[$0]) } == ["day"]) + } + + @Test func latinTermsBesideHanHighlightOriginalText() throws { + let store = try IndexStore(directory: folder.appendingPathComponent("mixed-scripts")) + try add(store, path: "/synthetic/mixed.txt", body: "学习Swift语言。") + let hit = try #require(try Searcher(store: store, loadVectors: false).search("swift", queryVector: nil).first) + #expect(hit.excerpt == "学习«Swift»语言。") + #expect(hit.excerpt?.contains("\u{E000}") == false) + } + + @Test func stemmingNeverShowsIndexedHanSeparators() throws { + let store = try IndexStore(directory: folder.appendingPathComponent("stemmed-han")) + try add(store, path: "/synthetic/stemmed.txt", body: "中文 running。") + let hit = try #require(try Searcher(store: store, loadVectors: false).search("runs", queryVector: nil).first) + #expect(hit.excerpt == "中文 running。") + let text = String(repeating: "甲", count: 1000) + "Swift" + String(repeating: "乙", count: 1000) + let snippet = try #require(SearchText.snippet(text, terms: ["swift"])) + #expect(snippet.count < 150) + #expect(snippet.contains("«Swift»")) + } + + @Test func chineseChunksRetainOverlapWithoutWhitespace() { + let text = String(repeating: "甲", count: 499) + "租房合同" + String(repeating: "乙", count: 1100) + "尾部标记" + let chunks = DocExtractor.chunks(text, size: 500, overlap: 100, limit: 100) + #expect(chunks.count > 2) + #expect(chunks.allSatisfy { $0.count <= 500 }) + #expect(chunks[1].prefix(100) == chunks[0].suffix(100)) + #expect(chunks.contains { $0.contains("租房合同") }) // phrase spans the first hard boundary. + #expect(chunks.last?.hasSuffix("尾部标记") == true) + } + + @Test func hanPhrasesDoNotBridgeSourceSeparatorsOrChooseTheWrongPage() throws { + let store = try IndexStore(directory: folder.appendingPathComponent("separators")) + let id = try add(store, path: "/synthetic/pages.pdf") + try store.addSegment(file: id, kind: .page, modality: .text, loc: 1, text: "租，房 租。房 租\n房") + let searcher = try Searcher(store: store, loadVectors: false) + #expect(try searcher.search("租房", queryVector: nil).isEmpty) + try store.addSegment(file: id, kind: .page, modality: .text, loc: 2, text: "租房合同已经签署") + let hits = try Searcher(store: store, loadVectors: false).search("租房", queryVector: nil) + #expect(hits.first?.loc == 2) + #expect(hits.first?.excerpt?.contains("«租房»") == true) + } + + @Test func schemaFiveKeepsStoredTextSearchableAfterUpgrade() throws { + let directory = folder.appendingPathComponent("migration") + var store: IndexStore? = try IndexStore(directory: directory) + let id = try add(store!, path: "/synthetic/年度合同备份.txt", body: "连续中文租房记录", inode: 1) + // Model the old unsplit word representation, retaining its original text. + try store!.db.run("UPDATE fts SET name = '年度合同备份.txt' WHERE rowid IN (SELECT id FROM segments WHERE kind = 'name')") + try store!.db.run("UPDATE fts SET body = '连续中文租房记录', original = NULL WHERE rowid IN (SELECT id FROM segments WHERE kind = 'chunk')") + try store!.finish(id, state: "done", detail: nil) + try store!.db.execute("PRAGMA user_version = 5") + store = nil + let reopened = try IndexStore(directory: directory) + let searcher = try Searcher(store: reopened, loadVectors: false) + #expect(try searcher.search("租房", queryVector: nil).map(\.path) == ["/synthetic/年度合同备份.txt"]) + #expect(try searcher.search("年度合同备份", queryVector: nil).map(\.path) == ["/synthetic/年度合同备份.txt"]) + } + + @Test func bodySegmentationLeavesFilenameSubstringFallbackCompatible() throws { + let store = try IndexStore(directory: folder.appendingPathComponent("name-compatibility")) + let id = try add(store, path: "/synthetic/年度合同备份.txt", body: "连续中文租房记录") + func namesContaining(_ needle: String) throws -> [Int64] { + var files: [Int64] = [] + // Same unspaced representation required by the existing filename-fallback PR. + try store.db.query(""" + SELECT s.file FROM fts JOIN segments s ON s.id = fts.rowid + WHERE s.kind = 'name' AND instr(fts.name, ?) > 0 + """, [.text(needle)]) { files.append($0.int(0)) } + return files + } + #expect(try namesContaining("合同") == [id]) + #expect(try Searcher(store: store, loadVectors: false).search("年度合同", queryVector: nil).first?.path + == "/synthetic/年度合同备份.txt") + try store.move(id, to: "/synthetic/中文目录/年度发票备份.txt") + #expect(try namesContaining("合同").isEmpty) + #expect(try namesContaining("发票") == [id]) + let hits = try Searcher(store: store, loadVectors: false).search("目录", queryVector: nil) + #expect(hits.first?.path == "/synthetic/中文目录/年度发票备份.txt") + #expect(hits.first?.excerpt?.contains("«目录»") == true) + } + + @Test func schemaFivePreservesOfflineCachedTextAndVectors() throws { + let directory = folder.appendingPathComponent("migration-budget") + var store: IndexStore? = try IndexStore(directory: directory) + let short = try add(store!, path: "/synthetic/short.txt", body: "租房合同", inode: 1) + let long = try add(store!, path: "/synthetic/long.txt", body: String(repeating: "甲", count: 1800), inode: 2) + let code = try store!.insert(Candidate(path: "/synthetic/code.md", kind: .code, size: 9000, + mtime: 1, inode: 3, device: 1)) + try store!.addSegment(file: code, kind: .lines, modality: .text, loc: 1, locEnd: 1, + text: String(repeating: "a", count: 4000), vector: FakeEmbedder.vector(for: "old")) + for id in [short, long, code] { try store!.finish(id, state: "done") } + try store!.db.execute("PRAGMA user_version = 5") + store = nil + let reopened = try IndexStore(directory: directory) + var states: [Int64: String] = [:] + try reopened.db.query("SELECT id, state FROM files") { states[$0.int(0)] = $0.text(1) } + #expect(states[short] == "done") + #expect(states[long] == "done") + #expect(states[code] == "done") + var version: Int64 = 0 + try reopened.db.query("PRAGMA user_version") { version = $0.int(0) } + #expect(version == 6) + var savedVectors = 0 + try reopened.db.query("SELECT count(*) FROM segments WHERE file = ? AND vec IS NOT NULL", [.int(code)]) { + savedVectors = Int($0.int(0)) + } + #expect(savedVectors == 1) + var savedBody = "" + try reopened.db.query("SELECT coalesce(fts.original, fts.body) FROM fts JOIN segments s ON s.id = fts.rowid WHERE s.file = ? AND s.kind = 'chunk'", [.int(long)]) { + savedBody = $0.text(0) ?? "" + } + #expect(savedBody == String(repeating: "甲", count: 1800)) + #expect(try Searcher(store: reopened, loadVectors: false).search("租房", queryVector: nil).first?.path + == "/synthetic/short.txt") + } + + @Test func longDocumentTailReachesTheEmbedder() throws { + let files = folder.appendingPathComponent("document-files") + try FileManager.default.createDirectory(at: files, withIntermediateDirectories: true) + let url = files.appendingPathComponent("long.txt") + let text = String(repeating: "公开合成的中文资料", count: 800) + "尾部检索标记" + try text.write(to: url, atomically: true, encoding: .utf8) + let store = try IndexStore(directory: folder.appendingPathComponent("document-index")) + let embedder = FakeEmbedder() + var options = IndexOptions() + options.docStep = 1 + let indexer = try Indexer(store: store, options: options, makeEmbedder: { embedder }) + _ = try indexer.sync([files]) + #expect(try indexer.run()["partial"] == 1) + #expect(!embedder.textInputs.contains { $0.contains("尾部检索标记") }) + #expect(try indexer.finishLongFiles()["done"] == 1) + #expect(embedder.textInputs.contains { $0.contains("尾部检索标记") }) + #expect(embedder.textInputs.allSatisfy { $0.utf8.count + 2 <= embedder.textTokenLimit }) + #expect(try Searcher(store: store, loadVectors: false).search("尾部检索标记", queryVector: nil).first?.path == url.path) + } + + @Test func impossibleFileBudgetDoesNotBlockTheQueue() throws { + let files = folder.appendingPathComponent("budget-files") + try FileManager.default.createDirectory(at: files, withIntermediateDirectories: true) + for name in [String(repeating: "title", count: 20), "short"] { + try "Public text.".write(to: files.appendingPathComponent(name + ".txt"), atomically: true, encoding: .utf8) + } + let store = try IndexStore(directory: folder.appendingPathComponent("failed-budget-index")) + let embedder = SmallBudgetEmbedder() + let indexer = try Indexer(store: store, options: IndexOptions(), makeEmbedder: { embedder }) + _ = try indexer.sync([files]) + let result = try indexer.run() + #expect(result["failed"] == 1) + #expect(result["done"] == 1) + #expect(try store.pendingCount() == 0) + } + + #if os(macOS) + @Test func pdfTailIsEmbeddedAndKeepsTheOriginalPageNumber() throws { + let files = folder.appendingPathComponent("files") + try FileManager.default.createDirectory(at: files, withIntermediateDirectories: true) + let url = files.appendingPathComponent("synthetic-long-page.pdf") + var box = CGRect(x: 0, y: 0, width: 1600, height: 6000) + let context = try #require(CGContext(url as CFURL, mediaBox: &box, nil)) + let long = String(repeating: "这是一段公开合成的中文资料。", count: 600) + "尾部检索标记" + for text in [String(repeating: "First synthetic page. ", count: 20), long] { + context.beginPDFPage(nil) + NSGraphicsContext.saveGraphicsState() + NSGraphicsContext.current = NSGraphicsContext(cgContext: context, flipped: false) + (text as NSString).draw(in: CGRect(x: 60, y: 60, width: 1480, height: 5800), + withAttributes: [.font: NSFont.systemFont(ofSize: 12)]) + NSGraphicsContext.restoreGraphicsState() + context.endPDFPage() + } + context.closePDF() + let store = try IndexStore(directory: folder.appendingPathComponent("pdf-index")) + let embedder = FakeEmbedder() + let indexer = try Indexer(store: store, options: IndexOptions(), makeEmbedder: { embedder }) + _ = try indexer.sync([files]) + #expect(try indexer.run()["done"] == 1) + #expect(embedder.textInputs.contains { $0.contains("尾部检索标记") }) + let hits = try Searcher(store: store, loadVectors: false).search("标记", queryVector: nil) + #expect(hits.first?.path == url.path) + #expect(hits.first?.loc == 2) + } + + #endif + +} + +private final class SmallBudgetEmbedder: Embedder { + private let base = FakeEmbedder() + var info: WorkerInfo { base.info } + var textTokenLimit: Int { 80 } + func embedTexts(_ texts: [String]) throws -> [[Float]] { try base.embedTexts(texts) } + func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] { try base.embedImages(paths, budget: budget) } + func embedAudio(_ paths: [String]) throws -> [Embedding] { try base.embedAudio(paths) } + func close() {} +} diff --git a/Tests/DigUpKitTests/CodeTests.swift b/Tests/DigUpKitTests/CodeTests.swift index f03cea13..4db23b96 100644 --- a/Tests/DigUpKitTests/CodeTests.swift +++ b/Tests/DigUpKitTests/CodeTests.swift @@ -72,10 +72,10 @@ import Testing #expect(pieces.last?.last == 102) } - @Test func aVeryLongLineIsCutForTheModel() { + @Test func aVeryLongLineIsPreservedForTokenizerSplitting() { let pieces = CodeExtractor.lines("let a = 1\n" + String(repeating: "y", count: 9000) + "\nlet b = 2\n", size: 1800) #expect(pieces.map(\.first) == [1, 2, 3]) - #expect(pieces[1].text.count == CodeExtractor.longestPiece) + #expect(pieces[1].text.count == 9000) } @Test func windowsLineEndingsAndBlankEdges() { @@ -228,6 +228,50 @@ import Testing // MARK: The code index + @Test func chineseCodePassagesRespectTokenBudgetAndKeepTheirLine() throws { + try write("chinese/.git/HEAD", "ref: refs/heads/main\n") + let longLine = String(repeating: "公开合成的中文资料", count: 600) + "尾部检索标记" + let url = try write("chinese/README.md", "# Public synthetic notes\n" + longLine + "\nFinal line.\n") + let store = try IndexStore(directory: folder.appendingPathComponent("budget-index"), file: IndexStore.codeFile) + let embedder = BudgetCheckingCodeEmbedder() + let indexer = try Indexer(store: store, options: codeOptions, makeEmbedder: { embedder }) + _ = try indexer.sync([folder.appendingPathComponent("chinese")]) + #expect(try indexer.run() == ["done": 1]) + #expect(embedder.inputs.count > 3) + #expect(embedder.inputs.allSatisfy { embedder.textTokenCount($0) <= embedder.textTokenLimit }) + #expect(embedder.inputs.contains { $0.contains("尾部检索标记") }) + let hits = try Searcher(store: store, loadVectors: false).search("尾部检索标记", queryVector: nil) + let hit = try #require(hits.first) + #expect(hit.path == url.path) + #expect(hit.segment == .lines && hit.loc == 2 && hit.locEnd == 2) + } + + @Test func chineseMultilineBudgetSplitKeepsOriginalFocusLine() throws { + try write("multiline/.git/HEAD", "ref: refs/heads/main\n") + // One character-sized code piece, but several tokenizer-sized parts. Repeated + // lines make searching for each part's text an ambiguous way to recover offsets. + let repeated = String(repeating: "公开合成的中文资料", count: 12) + let lines = ["# Public synthetic notes"] + Array(repeating: repeated, count: 12) + ["尾部检索标记"] + let text = lines.joined(separator: "\n") + "\n" + #expect(CodeExtractor.lines(text).count == 1) + let url = try write("multiline/README.md", text) + let store = try IndexStore(directory: folder.appendingPathComponent("multiline-index"), file: IndexStore.codeFile) + let embedder = BudgetCheckingCodeEmbedder() + let indexer = try Indexer(store: store, options: codeOptions, makeEmbedder: { embedder }) + _ = try indexer.sync([folder.appendingPathComponent("multiline")]) + #expect(try indexer.run() == ["done": 1]) + #expect(embedder.inputs.count > 1) + #expect(embedder.inputs.allSatisfy { embedder.textTokenCount($0) <= embedder.textTokenLimit }) + #expect(embedder.inputs.contains { $0.contains("尾部检索标记") }) + let hit = try #require(try Searcher(store: store, loadVectors: false) + .search("尾部检索标记", queryVector: nil).first) + #expect(hit.path == url.path && hit.segment == .lines) + #expect(hit.focus == 14) + let firstLine = try #require(hit.loc) + #expect(firstLine > 1) + #expect(hit.locEnd == 14) + } + @Test func codeIndexIsAFileOfItsOwn() throws { try write("billing/.git/HEAD", "ref: refs/heads/main\n") try write("billing/retry.py", [function("retryWithBackoff", body: 22), function("circuit", body: 22), @@ -339,3 +383,24 @@ import Testing #expect(estimate.bytes > 21 * 1600 && estimate.bytes < 200_000) } } + +/// Deterministic tokenizer stand-in that rejects overflow as both production runtimes do. +/// Kept local to code regressions so existing FakeEmbedder-based tests retain their behavior. +private final class BudgetCheckingCodeEmbedder: Embedder { + private let base = FakeEmbedder() + var info: WorkerInfo { base.info } + var textTokenLimit: Int { 2048 } + var inputs: [String] { base.textInputs } + func textTokenCount(_ text: String) -> Int { text.utf8.count + 2 } + func embedTexts(_ texts: [String]) throws -> [[Float]] { + guard texts.allSatisfy({ textTokenCount($0) <= textTokenLimit }) else { + throw WorkerError.failed("code regression: prompt exceeds tokenizer budget") + } + return try base.embedTexts(texts) + } + func embedImages(_ paths: [String], budget: Int) throws -> [Embedding] { + try base.embedImages(paths, budget: budget) + } + func embedAudio(_ paths: [String]) throws -> [Embedding] { try base.embedAudio(paths) } + func close() { base.close() } +} diff --git a/Tests/DigUpKitTests/FakeEmbedder.swift b/Tests/DigUpKitTests/FakeEmbedder.swift index 87771e17..43e80f90 100644 --- a/Tests/DigUpKitTests/FakeEmbedder.swift +++ b/Tests/DigUpKitTests/FakeEmbedder.swift @@ -6,6 +6,7 @@ final class FakeEmbedder: Embedder { let info = WorkerInfo(model: "fake", revision: "0000000", runtime: "test", dtype: "f32", dim: Vectors.dimension, textOnly: false, loadSeconds: nil) private(set) var texts = 0 + private(set) var textInputs: [String] = [] static func vector(for text: String) -> [Float] { var state = UInt64(5381) @@ -21,6 +22,7 @@ final class FakeEmbedder: Embedder { func embedTexts(_ texts: [String]) throws -> [[Float]] { self.texts += texts.count + textInputs += texts return texts.map(Self.vector(for:)) } diff --git a/Tests/DigUpKitTests/TextBudgetTests.swift b/Tests/DigUpKitTests/TextBudgetTests.swift new file mode 100644 index 00000000..83d873e2 --- /dev/null +++ b/Tests/DigUpKitTests/TextBudgetTests.swift @@ -0,0 +1,36 @@ +import Foundation +import Testing +@testable import DigUpKit + +@Suite struct TextBudgetTests { + @Test func budgetIncludesTitleAndKeepsOverlappingHanAndFinalContent() throws { + let text = String(repeating: "甲", count: 110) + "租房合同" + String(repeating: "乙", count: 300) + "尾部标记" + let title = "公开合成标题" + let count = { (input: String) in input.utf8.count + 2 } + let parts = try TextBudget.chunks(text, title: title, limit: 220, count: count) + #expect(parts.count > 3) + #expect(parts.allSatisfy { count(Prompts.document(title: title, text: $0)) <= 220 }) + #expect(parts.contains { $0.contains("租房合同") }) + #expect(parts.last?.hasSuffix("尾部标记") == true) + for pair in zip(parts, parts.dropFirst()) { + #expect((1...min(200, pair.0.count / 4)).contains { + pair.0.suffix($0) == pair.1.prefix($0) + }) + } + } + + @Test func smallEnglishAndBengaliArabicPassagesStayUnchanged() throws { + for text in ["A payment error appeared.", "সুন্দরবনে আমরা কুমিরও দেখেছি।", "قَالَ إِنَّ الْمَكْتَبَةَ تُغْلَقُ"] { + #expect(try TextBudget.chunks(text, title: "synthetic", limit: 2048, + count: { $0.utf8.count + 2 }) == [text]) + } + } + + @Test func impossibleTitleBudgetFailsInsteadOfDiscardingText() { + #expect(throws: (any Error).self) { + _ = try TextBudget.chunks("合同", title: String(repeating: "标题", count: 100), limit: 50, + count: { $0.utf8.count + 2 }) + } + } + +} diff --git a/Tests/DigUpKitTests/WorkerBudgetTests.swift b/Tests/DigUpKitTests/WorkerBudgetTests.swift new file mode 100644 index 00000000..c85d132a --- /dev/null +++ b/Tests/DigUpKitTests/WorkerBudgetTests.swift @@ -0,0 +1,56 @@ +#if !os(Windows) +import Foundation +import Testing +@testable import DigUpKit + +/// The JSON-lines contract, without a model or a Python environment. +@Suite struct WorkerBudgetTests { + private func withWorker(reply: String, limit: Int? = 77, _ test: (EmbeddingWorker) throws -> Void) throws { + let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: directory) } + let script = directory.appendingPathComponent("worker.sh") + let limitField = limit.map { ",\"text_token_limit\":\($0)" } ?? "" + let ready = "{\"ready\":true,\"info\":{\"model\":\"mock\",\"revision\":\"1\",\"runtime\":\"test\",\"dtype\":\"f32\",\"dim\":768,\"text_only\":true\(limitField)}}" + try """ + printf '%s\\n' '\(ready)' + while IFS= read -r line; do + case "$line" in + *token_count*) printf '%s\\n' '\(reply)' ;; + *quit*) exit 0 ;; + esac + done + """.write(to: script, atomically: true, encoding: .utf8) + let worker = try EmbeddingWorker(WorkerConfig(executable: URL(fileURLWithPath: "/bin/sh"), arguments: [script.path]), + textOnly: true) + defer { worker.close() } + try test(worker) + } + + @Test func countsAndSnakeCaseLimitRoundTrip() throws { + try withWorker(reply: #"{"id":1,"ok":true,"counts":[42]}"#) { worker in + #expect(worker.info.textOnly) + #expect(worker.info.textTokenLimit == 77) + #expect(worker.textTokenLimit == 77) + let count = try worker.textTokenCount("title: public | text: 租房") + #expect(count == 42) + } + } + + @Test func legacyUnsupportedCountFailsExplicitly() throws { + try withWorker(reply: #"{"id":1,"ok":false,"error":"unknown op token_count"}"#, limit: nil) { worker in + #expect(worker.textTokenLimit == 2048) + #expect(throws: WorkerError.self) { try worker.textTokenCount("public") } + } + } + + @Test(arguments: [#"{"id":2,"ok":true,"counts":[42]}"#, + #"{"id":1,"ok":true,"counts":[-1]}"#, + #"{"id":1,"ok":true,"counts":[1,2]}"#]) + func malformedCountsFailExplicitly(_ reply: String) throws { + try withWorker(reply: reply) { worker in + #expect(throws: WorkerError.self) { try worker.textTokenCount("public") } + } + } +} +#endif diff --git a/Tests/LlamaRuntimeTests/NativeTextBudgetTests.swift b/Tests/LlamaRuntimeTests/NativeTextBudgetTests.swift new file mode 100644 index 00000000..17d86845 --- /dev/null +++ b/Tests/LlamaRuntimeTests/NativeTextBudgetTests.swift @@ -0,0 +1,36 @@ +import DigUpKit +import Foundation +import Testing +@testable import LlamaRuntime + +@Suite(.serialized) struct NativeTextBudgetTests { + static let textModelPresent = (try? LlamaModels.locate()) != nil + + @Test(.enabled(if: textModelPresent)) func oversizedTextNeverSilentlyDropsItsTail() throws { + let embedder = try LlamaEmbedder(models: LlamaModels.locate(), textOnly: true) + defer { embedder.close() } + // Synthetic unspaced Han is far beyond the query encoder's 512-token context. + #expect(throws: (any Error).self) { + _ = try embedder.embedTexts([String(repeating: "租房合同", count: 1500) + "尾部标记"]) + } + } + + @Test(.enabled(if: textModelPresent)) func actualTokenizerBudgetPreservesEverySyntheticRecord() throws { + let embedder = try LlamaEmbedder(models: LlamaModels.locate(), textOnly: true) + defer { embedder.close() } + let records = (1...180).map { "公开记录\($0)的租房合同。" } + let parts = try embedder.documentChunks(records.joined() + "尾部标记", title: "公开合成回归") + #expect(parts.count > 1) + #expect(try parts.allSatisfy { + try embedder.textTokenCount(Prompts.document(title: "公开合成回归", text: $0)) <= embedder.textTokenLimit + }) + #expect(records.allSatisfy { record in parts.contains { $0.contains(record) } }) + #expect(parts.last?.hasSuffix("尾部标记") == true) + let vectors = try embedder.embedTexts(parts.prefix(2).map { + Prompts.document(title: "公开合成回归", text: $0) + }) + #expect(vectors.count == 2) + #expect(vectors.allSatisfy { $0.count == embedder.info.dim && $0.allSatisfy(\.isFinite) }) + print("NATIVE_TEXT_BUDGET actual tokenizer: \(parts.count) complete parts within \(embedder.textTokenLimit) tokens; 2 inference vectors validated") + } +} diff --git a/docs/han-body-search.md b/docs/han-body-search.md new file mode 100644 index 00000000..dbfcb113 --- /dev/null +++ b/docs/han-body-search.md @@ -0,0 +1,47 @@ +# Han body search and complete text embeddings + +This change is based on upstream `ec28973268b81f836be72187ff23b3a16ef4fd97`. + +## Scope + +- Index Han characters in document, PDF, OCR, code, and folder bodies as adjacent FTS tokens. Quoted query phrases preserve contiguity; punctuation and whitespace cannot disappear between characters. +- Keep `fts.name` normalized and **unspaced**, including after moves and migration. This preserves whole-name/prefix search and compatibility with the separate filename substring proposal (`instr(fts.name, ...)`). No filename trigram index is added. +- Highlight original text, including mixed Latin/Han text, without exposing indexing separators. Keep excerpts bounded for unspaced Han. +- Split complete title-plus-body prompts against the runtime tokenizer, including special tokens. Preserve document continuation offsets and PDF page/code line locations. Remove the old 4,000-character code-line cut and PDF text prefix cut. Both runtimes reject oversized requests instead of silently truncating them. +- Forward token counts through the JSON-lines worker and the indexing service's shared-model wrapper. + +No app identity, index location, media recovery/ranking, startup, localization, packaging, or model-version changes are included. + +## Existing indexes + +Schema 5 becomes schema 6 in one transaction. Bodies are rebuilt from `coalesce(original, body)`; names, vectors, file state, and segment locations remain intact. Migration does not read source files or require them to be connected. Older binaries cannot open schema 6, as with other schema upgrades. + +An old vector may already have lost a long passage's tail. Migration cannot reconstruct missing source text or replace those embeddings without reading the file and loading the tokenizer. The new splitting applies to new/changed files. To repair unchanged files, explicitly reindex when their folders are available, for example: + +```sh +digup index /path/to/folder --index /path/to/index --reindex +``` + +`--reindex` resets the selected index, so include all folders that belong in it; it is not a per-file repair command. Use the corresponding code index with `--code` when needed. Keep an index backup if rollback to an older binary is required. + +## Regression coverage + +`ChineseSearchTests`, `TextBudgetTests`, `CodeTests`, and `WorkerBudgetTests` cover body phrase contiguity, original snippets, supplementary-plane Han, filename compatibility, schema-5/offline-cache preservation, overlap, tail coverage, multi-pass documents, >4,000-character code lines, source line focus, impossible per-file budgets, and token-count protocol validation. + +`NativeTextBudgetTests` uses the actual pinned model when available. It checks oversized-input rejection, every synthetic record and final marker surviving a tokenizer-budget split, and finite 768-dimensional inference vectors. It explicitly skips without the model, which is not an inference pass. Python worker contract tests use deterministic stubs and do not load MLX: + +```sh +python -m unittest discover -s worker -p test_text_budget.py -v +``` + +### Measured before/after (Linux, Swift 6.2) + +- Same nine public synthetic keyword-only cases: upstream R@1 **2/9**, this change **9/9**. This is a narrow lexical regression corpus, not a multilingual semantic ranking benchmark. +- Actual pinned text model, query encoder's 512-token budget: upstream accepts a 3,013-token input and two different tails produce identical vectors. This change returns an explicit over-budget error. +- Actual tokenizer: 180 synthetic Han records plus a final marker become five complete parts within 512 tokens; every record survives and two inference vectors are validated. +- Full upstream suite: 92 tests, seven assertion issues. Full changed suite: 113 tests, the same seven baseline issues, no additional failures. Six are portable TIFF/HEIC/WebP orientation assertions with this environment's image reader; one is `placeInItsRepo` encountering the sandbox's synthetic `.git` parent. The final run excluding those two affected test functions passes all **111 selected tests with no skips**, including the actual-model budget tests and existing media/text `ReferenceVectorTests` with both pinned models. The long-document reference passes; no reference vectors or model version were changed. +- Python tokenizer contract: three tests pass. + +For Linux native model runs, set `DIGUP_MODELS` to the model folder and `DIGUP_BACKENDS` to `Vendor/llama/lib`; a development test binary does not carry ggml backend modules beside itself. All model files are pinned by `ModelFiles` and checksum-verified. + +The upstream `evals/` semantic suites need their generated `DigUpTestbed` source corpus, which is not bundled here. They are not claimed as run. macOS PDF extraction/rendering, Apple frameworks, and packaged-app validation remain separate from Linux engine validation. diff --git a/worker/digup_worker.py b/worker/digup_worker.py index 30158839..fa98293e 100644 --- a/worker/digup_worker.py +++ b/worker/digup_worker.py @@ -100,16 +100,25 @@ def __init__(self, model, text_only): "dim": DIM, "text_only": text_only, "load_seconds": round(time.perf_counter() - started, 2), + "text_token_limit": 512 if text_only else 2048, } def _vectors(self, output): return np.array(output.text_embeds.astype(self.mx.float32)) + def token_counts(self, inputs): + if not inputs: + return [] + tokens = self.tokenizer(inputs, padding=False, truncation=False, add_special_tokens=True) + return [len(ids) for ids in tokens["input_ids"]] + def text(self, inputs): + if any(count > self.info["text_token_limit"] for count in self.token_counts(inputs)): + raise ValueError("text exceeds the tokenizer budget; split it before embedding") rows = [] for i in range(0, len(inputs), BATCH["text"]): batch = inputs[i : i + BATCH["text"]] - tokens = self.tokenizer(batch, padding=True, truncation=True, max_length=2048, return_tensors="np") + tokens = self.tokenizer(batch, padding=True, truncation=False, return_tensors="np") rows.append(self._vectors(self.model(**{k: self.mx.array(v) for k, v in tokens.items()}))) return np.concatenate(rows), [None] * len(inputs) @@ -178,6 +187,9 @@ def main(): reply({"id": request_id, "ok": True, "info": embedder.info}) continue inputs = request.get("inputs") or [] + if op == "token_count": + reply({"id": request_id, "ok": True, "counts": embedder.token_counts(inputs)}) + continue if op == "text": vectors, errors = embedder.text(inputs) elif op in ("image", "audio"): diff --git a/worker/test_text_budget.py b/worker/test_text_budget.py new file mode 100644 index 00000000..5cb07e0d --- /dev/null +++ b/worker/test_text_budget.py @@ -0,0 +1,57 @@ +"""Tokenizer-budget contract without loading/downloading an MLX model.""" +import unittest +from types import SimpleNamespace + +import numpy as np + +from digup_worker import Embedder + + +class Tokenizer: + def __init__(self): + self.calls = [] + + def __call__(self, texts, **options): + self.calls.append(options) + rows = [[0] + list(text.encode('utf-8')) + [1] for text in texts] + if options.get('padding'): + width = max(map(len, rows)) + rows = [row + [0] * (width - len(row)) for row in rows] + return {'input_ids': rows} + + +class TextBudgetTests(unittest.TestCase): + def setUp(self): + self.embedder = Embedder.__new__(Embedder) + self.embedder.tokenizer = Tokenizer() + self.embedder.info = {'text_token_limit': 8} + self.embedder.mx = SimpleNamespace(array=np.array, float32=np.float32) + self.model_calls = [] + + def model(**tokens): + self.model_calls.append(tokens) + return SimpleNamespace(text_embeds=np.ones((len(tokens['input_ids']), 768), dtype=np.float32)) + + self.embedder.model = model + + def test_special_tokens_count_and_empty_batch(self): + self.assertEqual(self.embedder.token_counts(['abc', '租房']), [5, 8]) + self.assertEqual(self.embedder.token_counts([]), []) + self.assertTrue(self.embedder.tokenizer.calls[0]['add_special_tokens']) + self.assertFalse(self.embedder.tokenizer.calls[0]['truncation']) + + def test_exact_boundary_is_accepted_without_truncation(self): + vectors, errors = self.embedder.text(['租房']) + self.assertEqual(vectors.shape, (1, 768)) + self.assertEqual(errors, [None]) + self.assertEqual(self.model_calls[0]['input_ids'].shape, (1, 8)) + self.assertTrue(all(call['truncation'] is False for call in self.embedder.tokenizer.calls)) + + def test_overflow_rejected_before_any_model_call(self): + with self.assertRaisesRegex(ValueError, 'tokenizer budget'): + self.embedder.text(['ok', '租房x']) + self.assertEqual(self.model_calls, []) + + +if __name__ == '__main__': + unittest.main()